From 2fb87bfcf4dc9575371f235d0c12560bf375350f Mon Sep 17 00:00:00 2001 From: Shiju Date: Tue, 13 Oct 2026 01:20:12 +0530 Subject: [PATCH 1/5] fix(cli): backpressure sandbox uploads Stream archive uploads over a bounded SSH subsystem whose credit follows extraction stdin progress. Preserve the existing input budget and tar destination semantics, validate completion, and clean up cancelled transfers. Add protocol, lifecycle, and large-upload regression coverage without changing russh. Signed-off-by: Shiju --- Cargo.lock | 1 + crates/openshell-cli/Cargo.toml | 1 + crates/openshell-cli/src/ssh.rs | 179 ++-- crates/openshell-cli/src/ssh/upload.rs | 761 ++++++++++++++++++ crates/openshell-core/src/lib.rs | 1 + crates/openshell-core/src/upload.rs | 509 ++++++++++++ crates/openshell-supervisor-process/README.md | 31 + .../openshell-supervisor-process/src/ssh.rs | 89 +- .../src/ssh/input.rs | 167 +++- .../src/ssh/upload.rs | 586 ++++++++++++++ .../src/ssh/upload_tests.rs | 574 +++++++++++++ docs/how-it-works/sandboxes/overview.mdx | 10 + e2e/rust/e2e-podman.sh | 1 + e2e/rust/tests/upload_backpressure.rs | 121 +++ skills/openshell-cli/SKILL.md | 2 + 15 files changed, 2895 insertions(+), 138 deletions(-) create mode 100644 crates/openshell-cli/src/ssh/upload.rs create mode 100644 crates/openshell-core/src/upload.rs create mode 100644 crates/openshell-supervisor-process/README.md create mode 100644 crates/openshell-supervisor-process/src/ssh/upload.rs create mode 100644 crates/openshell-supervisor-process/src/ssh/upload_tests.rs create mode 100644 e2e/rust/tests/upload_backpressure.rs diff --git a/Cargo.lock b/Cargo.lock index 6a2c214e2c..b44a5c2ad9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4280,6 +4280,7 @@ dependencies = [ "rustls-pemfile", "serde", "serde_json", + "sha2 0.10.9", "tar", "temp-env", "tempfile", diff --git a/crates/openshell-cli/Cargo.toml b/crates/openshell-cli/Cargo.toml index 5c92eefe5b..a33a665201 100644 --- a/crates/openshell-cli/Cargo.toml +++ b/crates/openshell-cli/Cargo.toml @@ -65,6 +65,7 @@ anyhow = { workspace = true } # File archiving (tar-over-SSH sync) tar = "0.4" tempfile = "3" +sha2 = { workspace = true } # OIDC/Auth oauth2 = { version = "5", default-features = false, features = ["reqwest"] } diff --git a/crates/openshell-cli/src/ssh.rs b/crates/openshell-cli/src/ssh.rs index 30f7146d32..5de7a0686f 100644 --- a/crates/openshell-cli/src/ssh.rs +++ b/crates/openshell-cli/src/ssh.rs @@ -3,6 +3,8 @@ //! SSH connection and proxy utilities. +mod upload; + use crate::color::Colorize; use crate::tls::{ GrpcClient, TlsOptions, grpc_client, owner_replica, retry_unrouted, routed_to, with_owner_hint, @@ -19,6 +21,9 @@ use openshell_core::proto::{ CreateSshSessionRequest, GetSandboxRequest, SandboxPhase, SshRelayTarget, TcpForwardFrame, TcpForwardInit, tcp_forward_init, }; +use openshell_core::upload::{ + DestinationKind as UploadDestinationKind, ExpectedDestination, UploadPlan, +}; use std::fs; use std::future::Future; use std::io::{IsTerminal, Write}; @@ -64,7 +69,9 @@ const CONNECT_CHILD_TERMINATION_TIMEOUT: Duration = Duration::from_secs(5); const SSH_TRANSPORT_FAILURE_EXIT_CODE: i32 = 255; const SYNC_RETRY_ATTEMPTS: usize = 4; const SYNC_RETRY_DELAY: Duration = Duration::from_secs(2); -const UPLOAD_DESTINATION_CHANGED_EXIT_CODE: i32 = 73; +#[cfg(test)] +const UPLOAD_DESTINATION_CHANGED_EXIT_CODE: i32 = + openshell_core::upload::DESTINATION_CHANGED_EXIT_CODE; static DEFAULT_SSH_LOG_LEVEL: OnceLock<&'static str> = OnceLock::new(); /// Set the SSH log level derived from CLI verbosity before starting commands. @@ -1197,74 +1204,13 @@ fn local_upload_path_is_file_like(path: &Path) -> bool { }) } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum UploadDestinationKind { - Directory, - File, - Missing, -} - -impl UploadDestinationKind { - fn as_str(self) -> &'static str { - match self { - Self::Directory => "directory", - Self::File => "file", - Self::Missing => "missing", - } - } -} - -// Check every existing parent before interpreting a failed existence test as -// absence. An inaccessible directory, a file used as a parent, or a dangling -// parent symlink must fail inspection instead of selecting filename semantics. -// Only fixed tokens go to stdout; diagnostics remain on stderr. -const UPLOAD_DESTINATION_KIND_SH: &str = r#"upload_destination_kind() { - upload_path=$1 - upload_parent=$2 - while :; do - if [ -d "$upload_parent" ]; then - if [ ! -x "$upload_parent" ]; then - printf 'cannot inspect upload destination parent: %s\n' "$upload_parent" >&2 - return 1 - fi - elif [ -e "$upload_parent" ] || [ -L "$upload_parent" ]; then - printf 'upload destination parent is not a directory: %s\n' "$upload_parent" >&2 - return 1 - fi - case "$upload_parent" in - /|.) break ;; - */*) upload_parent=${upload_parent%/*}; [ -n "$upload_parent" ] || upload_parent=/ ;; - *) upload_parent=. ;; - esac - done - if [ -d "$upload_path" ]; then - printf directory - elif [ -e "$upload_path" ] || [ -L "$upload_path" ]; then - printf file - else - printf missing - fi -}"#; - -fn upload_destination_kind_invocation(path: &str) -> String { - let (parent, _) = split_sandbox_path(path); - format!( - "upload_destination_kind {} {}", - shell_escape(path), - shell_escape(parent), - ) -} - fn file_upload_destination_needs_probe(path: &str) -> bool { // Keep the existing /name heuristic: a top-level path means a directory. !path.ends_with('/') && split_sandbox_path(path).0 != "/" } fn build_upload_destination_probe_cmd(path: &str) -> String { - format!( - "{UPLOAD_DESTINATION_KIND_SH}\n{}", - upload_destination_kind_invocation(path), - ) + openshell_core::upload::build_destination_probe(path) } fn parse_upload_destination_kind(output: &str) -> Result { @@ -1293,32 +1239,19 @@ fn resolve_file_upload_destination<'a>( } } +#[cfg(test)] fn build_upload_extract_cmd( dest_dir: &str, expected_destination: Option<(&str, UploadDestinationKind)>, ) -> String { - let escaped_dest = shell_escape(dest_dir); - let extract = format!("mkdir -p -- {escaped_dest} && cat | tar xf - -C {escaped_dest}"); - let Some((path, kind)) = expected_destination else { - return extract; - }; - - // Recheck before mkdir or tar can replace a path whose type changed after - // planning. This detects changes, but is not atomic with extraction and - // does not pin a symlink's target against concurrent replacement. - format!( - "{UPLOAD_DESTINATION_KIND_SH}\n\ - upload_kind=$({probe}) || exit 1\n\ - if [ \"$upload_kind\" != {expected} ]; then\n\ - printf 'upload destination changed type before extraction: %s\\n' {path} >&2\n\ - exit {changed_status}\n\ - fi\n\ - {extract}", - probe = upload_destination_kind_invocation(path), - expected = kind.as_str(), - path = shell_escape(path), - changed_status = UPLOAD_DESTINATION_CHANGED_EXIT_CODE, - ) + openshell_core::upload::build_extract_command(&UploadPlan { + destination_directory: dest_dir.to_string(), + expected_destination: expected_destination.map(|(path, kind)| ExpectedDestination { + path: path.to_string(), + kind, + }), + }) + .expect("test upload destination is valid") } /// Stream an archive into the sandbox and return its final destination. @@ -1346,8 +1279,16 @@ async fn ssh_tar_upload( { let output = ssh_run_capture_stdout(&session, &build_upload_destination_probe_cmd(path)) .await + .map_err(|error| { + if sync_error_is_retryable(&error) { + error + } else { + Report::new(upload::UploadError::Remote(error.to_string())) + } + }) .wrap_err_with(|| format!("failed to inspect upload destination '{path}'"))?; - let kind = parse_upload_destination_kind(&output)?; + let kind = parse_upload_destination_kind(&output) + .map_err(|error| Report::new(upload::UploadError::Protocol(error.to_string())))?; let (directory, basename) = resolve_file_upload_destination(path, tar_name, kind); dest_dir = directory; *tar_name = basename; @@ -1370,40 +1311,21 @@ async fn ssh_tar_upload( ssh.arg("-T") .arg("-o") .arg("RequestTTY=no") + .arg("-s") .arg("sandbox") - .arg(build_upload_extract_cmd(dest_dir, expected_destination)) - .stdin(Stdio::piped()) - .stdout(Stdio::null()) - .stderr(Stdio::inherit()); - - let mut child = ssh.spawn().into_diagnostic()?; - let stdin = child - .stdin - .take() - .ok_or_else(|| miette::miette!("failed to open stdin for ssh process"))?; - - // Wait even when the writer gets a broken pipe. A failed remote guard can - // close stdin early; its status must not become a retryable write error. - let archive_result = - tokio::task::spawn_blocking(move || -> Result<()> { write_upload_archive(stdin, source) }) - .await - .into_diagnostic()?; - let status = tokio::task::spawn_blocking(move || child.wait()) - .await - .into_diagnostic()? - .into_diagnostic()?; - - if status.code() == Some(UPLOAD_DESTINATION_CHANGED_EXIT_CODE) { - return Err(miette::miette!( - "upload destination changed type before extraction; retry the upload" - )); - } - if !status.success() { - return Err(miette::miette!( - "ssh tar extract exited with status {status}" - )); - } - archive_result?; + .arg(openshell_core::upload::SUBSYSTEM); + let plan = UploadPlan { + destination_directory: dest_dir.to_string(), + expected_destination: expected_destination.map(|(path, kind)| ExpectedDestination { + path: path.to_string(), + kind, + }), + }; + upload::run(ssh, plan, move |writer| { + write_upload_archive(writer, source) + }) + .await + .map_err(Report::new)?; Ok(final_path) } @@ -1815,6 +1737,9 @@ where } fn sync_error_is_retryable(error: &Report) -> bool { + if let Some(error) = error.downcast_ref::() { + return error.is_retryable(); + } let message = format!("{error:?}").to_ascii_lowercase(); [ "broken pipe", @@ -2557,6 +2482,22 @@ mod tests { assert!(!sync_error_is_retryable(&error)); } + #[test] + fn upload_errors_retry_only_explicit_transport_failures() { + for error in [ + upload::UploadError::Unsupported, + upload::UploadError::Protocol("connection reset in malformed frame".into()), + upload::UploadError::Remote("destination /connection/unavailable".into()), + upload::UploadError::Archive("broken pipe while reading source".into()), + upload::UploadError::Cancelled, + ] { + assert!(!sync_error_is_retryable(&Report::new(error))); + } + assert!(sync_error_is_retryable(&Report::new( + upload::UploadError::Transport("SSH relay disconnected".into()) + ))); + } + #[test] fn connect_recovery_starts_only_for_established_transport_failures() { assert!(should_recover_connect( diff --git a/crates/openshell-cli/src/ssh/upload.rs b/crates/openshell-cli/src/ssh/upload.rs new file mode 100644 index 0000000000..247e4e09c9 --- /dev/null +++ b/crates/openshell-cli/src/ssh/upload.rs @@ -0,0 +1,761 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Bounded tar upload over the supervisor's credited SSH subsystem. + +use std::io::{self, Write}; +use std::process::{Command, ExitStatus, Stdio}; +use std::time::Duration; + +use miette::Diagnostic; +use openshell_core::upload::{ + Frame, MAX_DATA_BYTES, MAX_DIAGNOSTIC_BYTES, ProtocolError, UploadPlan, UploadSummary, + WINDOW_BYTES, read_frame, write_frame, +}; +use sha2::{Digest, Sha256}; +use thiserror::Error; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::process::{Child, Command as TokioCommand}; +use tokio::sync::mpsc; +use tokio::task::AbortHandle; + +const NEGOTIATION_TIMEOUT: Duration = Duration::from_secs(10); +const CHILD_EXIT_TIMEOUT: Duration = Duration::from_secs(5); +/// The producer can hold at most two chunks in addition to its current buffer +/// and the chunk held by the driver. It cannot build an entire archive in RAM. +const PRODUCER_QUEUE_CHUNKS: usize = 2; + +#[derive(Debug, Error, Diagnostic)] +pub(super) enum UploadError { + #[error("upload SSH transport failed: {0}")] + Transport(String), + #[error( + "this sandbox supervisor does not support openshell-upload-v1; upgrade the sandbox supervisor before uploading" + )] + Unsupported, + #[error("upload protocol failed: {0}")] + Protocol(String), + #[error("sandbox upload failed: {0}")] + Remote(String), + #[error("failed to create upload archive: {0}")] + Archive(String), + #[error("upload cancelled")] + Cancelled, +} + +impl UploadError { + pub(super) fn is_retryable(&self) -> bool { + matches!(self, Self::Transport(_)) + } +} + +/// Own the SSH child through cancellation as well as ordinary error paths. +/// Tokio's kill-on-drop handles termination; the detached waiter also reaps it +/// when an enclosing command future is dropped by its caller. +struct ChildGuard(Option); + +impl ChildGuard { + fn child(&mut self) -> &mut Child { + self.0.as_mut().expect("SSH child is still owned") + } + + async fn finish(&mut self, terminate: bool) -> io::Result { + let child = self.child(); + if terminate { + child.start_kill()?; + } + let status = child.wait().await?; + self.0.take(); + Ok(status) + } +} + +impl Drop for ChildGuard { + fn drop(&mut self) { + if let Some(mut child) = self.0.take() { + let _ = child.start_kill(); + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + runtime.spawn(async move { + let _ = child.wait().await; + }); + } + } + } +} + +struct AbortTasks(Vec); + +impl Drop for AbortTasks { + fn drop(&mut self) { + for task in &self.0 { + task.abort(); + } + } +} + +/// Start the producer only after the subsystem has accepted the destination. +/// The archive callback uses the existing tar builder unchanged. +pub(super) async fn run

( + mut command: Command, + plan: UploadPlan, + producer: P, +) -> Result<(), UploadError> +where + P: FnOnce(&mut dyn Write) -> miette::Result<()> + Send + 'static, +{ + plan.validate() + .map_err(|error| UploadError::Protocol(error.to_string()))?; + command + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + let child = TokioCommand::from(command) + .kill_on_drop(true) + .spawn() + .map_err(|error| UploadError::Transport(error.to_string()))?; + let mut child = ChildGuard(Some(child)); + let stdin = child.child().stdin.take().expect("piped SSH stdin"); + let stdout = child.child().stdout.take().expect("piped SSH stdout"); + let stderr = child.child().stderr.take().expect("piped SSH stderr"); + let diagnostics = tokio::spawn(read_diagnostic(stderr)); + let diagnostic_guard = AbortTasks(vec![diagnostics.abort_handle()]); + + #[cfg(unix)] + let mut signals = super::TerminationSignals::new() + .map_err(|error| UploadError::Transport(error.to_string()))?; + #[cfg(unix)] + let result = tokio::select! { + result = transfer(stdout, stdin, plan, producer) => result, + _ = signals.recv() => Err(UploadError::Cancelled), + }; + #[cfg(not(unix))] + let result = transfer(stdout, stdin, plan, producer).await; + + let status = if result.is_err() { + child.finish(true).await + } else if let Ok(status) = tokio::time::timeout(CHILD_EXIT_TIMEOUT, child.finish(false)).await { + status + } else { + let _ = child.finish(true).await; + return Err(UploadError::Protocol( + "SSH did not exit after the upload result".to_string(), + )); + }; + let diagnostic = tokio::time::timeout(CHILD_EXIT_TIMEOUT, diagnostics) + .await + .ok() + .and_then(Result::ok) + .unwrap_or_default(); + drop(diagnostic_guard); + // OpenSSH reports subsystem rejection as status 255, like other SSH + // failures. Its explicit request diagnostic distinguishes an old server + // from a transient connection failure, without sending any archive bytes. + if matches!(result, Err(UploadError::Transport(_))) + && diagnostic.contains("subsystem request failed") + { + return Err(UploadError::Unsupported); + } + match result { + Err(UploadError::Transport(message)) if !diagnostic.trim().is_empty() => { + return Err(UploadError::Transport(format!("{message}: {diagnostic}"))); + } + result => result?, + } + let status = status.map_err(|error| UploadError::Transport(error.to_string()))?; + if !status.success() { + if status.code() == Some(super::SSH_TRANSPORT_FAILURE_EXIT_CODE) { + return Err(UploadError::Transport(format!("{status}: {diagnostic}"))); + } + return Err(UploadError::Protocol(format!( + "SSH exited with {status} after a successful result: {diagnostic}" + ))); + } + Ok(()) +} + +async fn read_diagnostic(mut reader: R) -> String { + let mut diagnostic = Vec::new(); + let mut chunk = [0; 4096]; + loop { + match reader.read(&mut chunk).await { + Ok(0) | Err(_) => break, + Ok(bytes) => { + diagnostic.extend_from_slice(&chunk[..bytes]); + if diagnostic.len() > MAX_DIAGNOSTIC_BYTES { + diagnostic.drain(..diagnostic.len() - MAX_DIAGNOSTIC_BYTES); + } + } + } + } + String::from_utf8_lossy(&diagnostic).into_owned() +} + +fn protocol_error(error: ProtocolError) -> UploadError { + // ProtocolError preserves I/O errors separately from invalid frames. + match error { + ProtocolError::Io(error) => UploadError::Transport(error.to_string()), + error => UploadError::Protocol(error.to_string()), + } +} + +async fn negotiate( + reader: &mut R, + writer: &mut W, + plan: UploadPlan, +) -> Result<(u64, usize), UploadError> { + let hello = read_frame(reader).await.map_err(protocol_error)?; + let Frame::Hello { + window_bytes, + max_data_bytes, + } = hello + else { + return Err(server_failure(hello, "expected HELLO")); + }; + // A peer cannot increase local memory use by advertising larger limits. + if window_bytes == 0 + || window_bytes > WINDOW_BYTES + || max_data_bytes == 0 + || usize::try_from(max_data_bytes).unwrap_or(usize::MAX) > MAX_DATA_BYTES + || max_data_bytes > window_bytes + { + return Err(UploadError::Protocol( + "supervisor advertised invalid upload limits".to_string(), + )); + } + write_frame(writer, &Frame::Begin(plan)) + .await + .map_err(protocol_error)?; + match read_frame(reader).await.map_err(protocol_error)? { + Frame::Ready => Ok((u64::from(window_bytes), max_data_bytes as usize)), + frame => Err(server_failure(frame, "expected READY")), + } +} + +fn server_failure(frame: Frame, expected: &str) -> UploadError { + match frame { + Frame::Error(failure) => { + UploadError::Remote(format!("{:?}: {}", failure.code, failure.message)) + } + _ => UploadError::Protocol(expected.to_string()), + } +} + +enum ProducerEvent { + Data(Vec), + Finished(Result<(), UploadError>), +} + +struct ChunkWriter { + sender: mpsc::Sender, + buffer: Vec, + max_chunk: usize, +} + +impl ChunkWriter { + fn flush_chunk(&mut self) -> io::Result<()> { + if !self.buffer.is_empty() { + let buffer = std::mem::replace(&mut self.buffer, Vec::with_capacity(self.max_chunk)); + self.sender + .blocking_send(ProducerEvent::Data(buffer)) + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "upload stopped"))?; + } + Ok(()) + } +} + +impl Write for ChunkWriter { + fn write(&mut self, mut bytes: &[u8]) -> io::Result { + let original_len = bytes.len(); + while !bytes.is_empty() { + let length = bytes.len().min(self.max_chunk - self.buffer.len()); + self.buffer.extend_from_slice(&bytes[..length]); + bytes = &bytes[length..]; + if self.buffer.len() == self.max_chunk { + self.flush_chunk()?; + } + } + Ok(original_len) + } + + fn flush(&mut self) -> io::Result<()> { + self.flush_chunk() + } +} + +fn spawn_producer

(producer: P, max_chunk: usize) -> mpsc::Receiver +where + P: FnOnce(&mut dyn Write) -> miette::Result<()> + Send + 'static, +{ + let (sender, receiver) = mpsc::channel(PRODUCER_QUEUE_CHUNKS); + tokio::task::spawn_blocking(move || { + let mut writer = ChunkWriter { + sender: sender.clone(), + buffer: Vec::with_capacity(max_chunk), + max_chunk, + }; + let result = producer(&mut writer) + .map_err(|error| UploadError::Archive(error.to_string())) + .and_then(|()| { + writer + .flush() + .map_err(|error| UploadError::Archive(error.to_string())) + }); + let _ = sender.blocking_send(ProducerEvent::Finished(result)); + }); + receiver +} + +/// Dedicated readers and writers let remote failure interrupt a blocked local +/// SSH pipe write. The frame reader is never cancelled between partial reads. +async fn transfer( + mut reader: R, + mut writer: W, + plan: UploadPlan, + producer: P, +) -> Result<(), UploadError> +where + R: AsyncRead + Unpin + Send + 'static, + W: AsyncWrite + Unpin + Send + 'static, + P: FnOnce(&mut dyn Write) -> miette::Result<()> + Send + 'static, +{ + let (window, max_chunk) = tokio::time::timeout( + NEGOTIATION_TIMEOUT, + negotiate(&mut reader, &mut writer, plan), + ) + .await + .map_err(|_| UploadError::Transport("upload negotiation timed out".to_string()))??; + let mut events = spawn_producer(producer, max_chunk); + let (incoming, mut frames) = mpsc::channel(4); + let reader_task = tokio::spawn(async move { + loop { + let frame = read_frame(&mut reader).await.map_err(protocol_error); + let terminal = frame.is_err(); + if incoming.send(frame).await.is_err() || terminal { + break; + } + } + }); + let (outgoing, mut queued) = mpsc::channel::(1); + let (completed, mut completions) = mpsc::channel(1); + let writer_task = tokio::spawn(async move { + while let Some(frame) = queued.recv().await { + let is_end = matches!(frame, Frame::End(_)); + let result = write_frame(&mut writer, &frame) + .await + .map_err(protocol_error); + let result = if result.is_ok() && is_end { + writer + .shutdown() + .await + .map_err(|error| UploadError::Transport(error.to_string())) + } else { + result + }; + let terminal = result.is_err() || is_end; + if completed.send(result).await.is_err() || terminal { + break; + } + } + }); + let _tasks = AbortTasks(vec![reader_task.abort_handle(), writer_task.abort_handle()]); + let mut sent = 0_u64; + let mut acknowledged = 0_u64; + let mut digest = Sha256::new(); + let mut pending = None; + let mut producer_finished = false; + let mut writing = false; + let mut end = None; + loop { + if !writing && end.is_none() { + if pending.as_ref().is_some_and(|data: &Vec| { + (sent - acknowledged) + .checked_add(data.len() as u64) + .is_some_and(|pending| pending <= window) + }) { + let data = pending.take().expect("pending upload data"); + sent = sent.checked_add(data.len() as u64).ok_or_else(|| { + UploadError::Protocol("upload byte count overflow".to_string()) + })?; + digest.update(&data); + outgoing + .send(Frame::Data(data)) + .await + .map_err(|_| UploadError::Transport("upload writer stopped".to_string()))?; + writing = true; + } else if producer_finished && pending.is_none() { + let summary = UploadSummary { + bytes: sent, + sha256: digest.clone().finalize().into(), + }; + outgoing + .send(Frame::End(summary)) + .await + .map_err(|_| UploadError::Transport("upload writer stopped".to_string()))?; + end = Some(summary); + writing = true; + } + } + tokio::select! { + frame = frames.recv() => { + match frame.ok_or_else(|| UploadError::Transport("upload reader stopped".to_string()))?? { + Frame::Ack(bytes) => { + if bytes < acknowledged || bytes > sent { + return Err(UploadError::Protocol("invalid cumulative upload ACK".to_string())); + } + acknowledged = bytes; + } + Frame::Result(summary) => { + if end != Some(summary) { + return Err(UploadError::Protocol("upload result count or digest does not match END".to_string())); + } + return Ok(()); + } + frame => return Err(server_failure(frame, "unexpected frame during upload")), + } + } + result = completions.recv(), if writing => { + result.ok_or_else(|| UploadError::Transport("upload writer stopped".to_string()))??; + writing = false; + } + event = events.recv(), if pending.is_none() && !producer_finished => { + match event.ok_or_else(|| UploadError::Archive("archive producer stopped".to_string()))? { + ProducerEvent::Data(data) => pending = Some(data), + ProducerEvent::Finished(result) => { + result?; + producer_finished = true; + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use openshell_core::upload::{FailureCode, UploadFailure}; + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + use tokio::io::{DuplexStream, ReadHalf, WriteHalf}; + + fn plan() -> UploadPlan { + UploadPlan { + destination_directory: "/tmp/upload".to_string(), + expected_destination: None, + } + } + + async fn accept(stream: DuplexStream) -> (ReadHalf, WriteHalf) { + let (mut read, mut write) = tokio::io::split(stream); + write_frame( + &mut write, + &Frame::Hello { + window_bytes: WINDOW_BYTES, + max_data_bytes: u32::try_from(MAX_DATA_BYTES).unwrap(), + }, + ) + .await + .unwrap(); + assert!(matches!( + read_frame(&mut read).await.unwrap(), + Frame::Begin(_) + )); + write_frame(&mut write, &Frame::Ready).await.unwrap(); + (read, write) + } + + fn payload(bytes: usize) -> Vec { + (0..bytes) + .map(|index| u8::try_from(index % 251).unwrap()) + .collect() + } + + async fn complete( + read: &mut ReadHalf, + write: &mut WriteHalf, + mut received: Vec, + ) -> Vec { + loop { + match read_frame(read).await.unwrap() { + Frame::Data(data) => { + assert!(data.len() <= MAX_DATA_BYTES); + received.extend_from_slice(&data); + write_frame(write, &Frame::Ack(received.len() as u64)) + .await + .unwrap(); + } + Frame::End(summary) => { + assert_eq!(summary.bytes, received.len() as u64); + assert_eq!(summary.sha256, <[u8; 32]>::from(Sha256::digest(&received))); + write_frame(write, &Frame::Result(summary)).await.unwrap(); + return received; + } + frame => panic!("unexpected client frame: {frame:?}"), + } + } + } + + #[tokio::test] + async fn stalled_consumer_blocks_at_credit_limit_then_resumes_with_exact_data() { + let expected = payload(6 * WINDOW_BYTES as usize + 123); + let bytes = expected.clone(); + let (client, server) = tokio::io::duplex(4096); + let (read, write) = tokio::io::split(client); + let client = tokio::spawn(transfer(read, write, plan(), move |writer| { + writer.write_all(&bytes).map_err(miette::Report::from_err)?; + Ok(()) + })); + let (mut read, mut write) = accept(server).await; + let mut received = Vec::new(); + while received.len() < WINDOW_BYTES as usize { + match read_frame(&mut read).await.unwrap() { + Frame::Data(data) => received.extend_from_slice(&data), + frame => panic!("unexpected client frame: {frame:?}"), + } + } + assert_eq!(received.len(), WINDOW_BYTES as usize); + assert!( + tokio::time::timeout(Duration::from_millis(50), read_frame(&mut read)) + .await + .is_err(), + "client exceeded credit while the extractor was stalled" + ); + assert!(!client.is_finished()); + write_frame(&mut write, &Frame::Ack(received.len() as u64)) + .await + .unwrap(); + assert_eq!(complete(&mut read, &mut write, received).await, expected); + client.await.unwrap().unwrap(); + } + + #[tokio::test] + async fn cancelling_stalled_transfer_stops_blocking_producer_and_closes_pipes() { + struct ProducerExit(Arc); + impl Drop for ProducerExit { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + let finished = Arc::new(AtomicBool::new(false)); + let producer_finished = finished.clone(); + let (client, server) = tokio::io::duplex(4096); + let (read, write) = tokio::io::split(client); + let client = tokio::spawn(transfer(read, write, plan(), move |writer| { + let _finished = ProducerExit(producer_finished); + writer + .write_all(&payload(8 * WINDOW_BYTES as usize)) + .map_err(miette::Report::from_err)?; + Ok(()) + })); + let (mut read, _write) = accept(server).await; + let mut received = 0; + while received < WINDOW_BYTES as usize { + let Frame::Data(data) = read_frame(&mut read).await.unwrap() else { + panic!("expected data before credit is exhausted"); + }; + received += data.len(); + } + assert_eq!(received, WINDOW_BYTES as usize); + assert!(!finished.load(Ordering::SeqCst)); + client.abort(); + assert!(client.await.unwrap_err().is_cancelled()); + tokio::time::timeout(Duration::from_secs(2), async { + while !finished.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled transfer must stop its blocking archive producer"); + let mut byte = [0]; + assert_eq!( + tokio::time::timeout(Duration::from_secs(2), read.read(&mut byte)) + .await + .expect("cancelled transfer must release its protocol pipes") + .unwrap(), + 0 + ); + } + + #[tokio::test] + async fn impossible_ack_is_rejected_without_retry() { + let (client, server) = tokio::io::duplex(4096); + let (read, write) = tokio::io::split(client); + let client = tokio::spawn(transfer(read, write, plan(), |writer| { + writer + .write_all(&payload(2 * WINDOW_BYTES as usize)) + .map_err(miette::Report::from_err)?; + Ok(()) + })); + let (_read, mut write) = accept(server).await; + write_frame(&mut write, &Frame::Ack(u64::MAX)) + .await + .unwrap(); + let error = client.await.unwrap().unwrap_err(); + assert!(matches!(error, UploadError::Protocol(_))); + assert!(!error.is_retryable()); + } + + #[tokio::test] + async fn regressing_ack_is_rejected() { + let (client, server) = tokio::io::duplex(4096); + let (read, write) = tokio::io::split(client); + let client = tokio::spawn(transfer(read, write, plan(), |writer| { + writer + .write_all(&payload(2 * WINDOW_BYTES as usize)) + .map_err(miette::Report::from_err)?; + Ok(()) + })); + let (mut read, mut write) = accept(server).await; + let Frame::Data(data) = read_frame(&mut read).await.unwrap() else { + panic!("expected data"); + }; + write_frame(&mut write, &Frame::Ack(data.len() as u64)) + .await + .unwrap(); + write_frame(&mut write, &Frame::Ack(data.len() as u64 - 1)) + .await + .unwrap(); + assert!(matches!( + client.await.unwrap().unwrap_err(), + UploadError::Protocol(_) + )); + } + + #[tokio::test] + async fn terminal_error_interrupts_blocked_pipe_write() { + let (client, server) = tokio::io::duplex(64); + let (read, write) = tokio::io::split(client); + let client = tokio::spawn(transfer(read, write, plan(), |writer| { + writer + .write_all(&payload(2 * WINDOW_BYTES as usize)) + .map_err(miette::Report::from_err)?; + Ok(()) + })); + let (_read, mut write) = accept(server).await; + write_frame( + &mut write, + &Frame::Error(UploadFailure { + code: FailureCode::InputWrite, + message: "extractor closed stdin: connection is unavailable".to_string(), + }), + ) + .await + .unwrap(); + let error = tokio::time::timeout(Duration::from_secs(1), client) + .await + .expect("remote failure must interrupt a blocked writer") + .unwrap() + .unwrap_err(); + assert!(matches!(error, UploadError::Remote(_))); + assert!(!error.is_retryable()); + } + + #[tokio::test] + async fn successful_exit_cannot_replace_terminal_result() { + let (client, server) = tokio::io::duplex(4096); + let (read, write) = tokio::io::split(client); + let client = tokio::spawn(transfer(read, write, plan(), |_| Ok(()))); + let (mut read, write) = accept(server).await; + assert!(matches!( + read_frame(&mut read).await.unwrap(), + Frame::End(_) + )); + drop(write); + drop(read); + assert!(client.await.unwrap().is_err()); + } + + #[tokio::test] + async fn mismatched_terminal_digest_is_rejected() { + let (client, server) = tokio::io::duplex(4096); + let (read, write) = tokio::io::split(client); + let client = tokio::spawn(transfer(read, write, plan(), |_| Ok(()))); + let (mut read, mut write) = accept(server).await; + let Frame::End(mut summary) = read_frame(&mut read).await.unwrap() else { + panic!("expected end"); + }; + summary.sha256[0] ^= 1; + write_frame(&mut write, &Frame::Result(summary)) + .await + .unwrap(); + assert!(matches!( + client.await.unwrap().unwrap_err(), + UploadError::Protocol(_) + )); + } + + #[cfg(unix)] + #[tokio::test] + async fn unavailable_subsystem_never_starts_archive_producer() { + let started = Arc::new(AtomicBool::new(false)); + let producer_started = started.clone(); + let mut command = Command::new("sh"); + command.args([ + "-c", + "printf 'subsystem request failed on channel 0\\n' >&2; exit 255", + ]); + let error = run(command, plan(), move |_| { + producer_started.store(true, Ordering::SeqCst); + Ok(()) + }) + .await + .unwrap_err(); + assert!(matches!(error, UploadError::Unsupported)); + assert!(!started.load(Ordering::SeqCst)); + assert!(!error.is_retryable()); + } + + #[cfg(unix)] + #[tokio::test] + async fn early_ssh_exit_never_starts_archive_producer() { + let started = Arc::new(AtomicBool::new(false)); + let producer_started = started.clone(); + let mut command = Command::new("sh"); + command.args(["-c", "exit 255"]); + let error = run(command, plan(), move |_| { + producer_started.store(true, Ordering::SeqCst); + Ok(()) + }) + .await + .unwrap_err(); + assert!(matches!(error, UploadError::Transport(_))); + assert!(!started.load(Ordering::SeqCst)); + assert!(error.is_retryable()); + } + + #[cfg(unix)] + #[tokio::test] + async fn dropping_upload_future_terminates_and_reaps_ssh_child() { + let temporary = tempfile::tempdir().unwrap(); + let pid_file = temporary.path().join("pid"); + let mut command = Command::new("sh"); + command.args([ + "-c", + "printf '%s' \"$$\" > \"$1\"; exec sleep 60", + "upload-test", + ]); + command.arg(&pid_file); + let upload = tokio::spawn(run(command, plan(), |_| Ok(()))); + let pid = tokio::time::timeout(Duration::from_secs(2), async { + loop { + if let Ok(pid) = tokio::fs::read_to_string(&pid_file).await + && let Ok(pid) = pid.parse::() + { + break pid; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + upload.abort(); + assert!(upload.await.unwrap_err().is_cancelled()); + tokio::time::timeout(Duration::from_secs(2), async { + while nix::sys::signal::kill(nix::unistd::Pid::from_raw(pid), None).is_ok() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("SSH child must be terminated and reaped on cancellation"); + } +} diff --git a/crates/openshell-core/src/lib.rs b/crates/openshell-core/src/lib.rs index d07f9241a4..d4bfb4579a 100644 --- a/crates/openshell-core/src/lib.rs +++ b/crates/openshell-core/src/lib.rs @@ -62,6 +62,7 @@ pub mod spiffe; pub mod telemetry; pub mod time; pub mod transport_errors; +pub mod upload; pub use config::{ AppArmorProfile, Config, GatewayAuthConfig, GatewayInterceptorBindingOverride, diff --git a/crates/openshell-core/src/upload.rs b/crates/openshell-core/src/upload.rs new file mode 100644 index 0000000000..c7969f9b06 --- /dev/null +++ b/crates/openshell-core/src/upload.rs @@ -0,0 +1,509 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Internal streaming-upload protocol shared by the CLI and SSH supervisor. +//! +//! SSH receive credit is replenished before application consumption. This +//! subsystem supplies its own bounded credit: an ACK means boundary stdin has +//! accepted bytes, while RESULT means extraction completed successfully. + +use serde::{Deserialize, Serialize}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; + +use crate::forward::shell_escape; + +pub const SUBSYSTEM: &str = "openshell-upload-v1"; +pub const WINDOW_BYTES: u32 = 1024 * 1024; +pub const MAX_DATA_BYTES: usize = 64 * 1024; +pub const MAX_METADATA_BYTES: usize = 16 * 1024; +pub const MAX_DIAGNOSTIC_BYTES: usize = 16 * 1024; +pub const DESTINATION_CHANGED_EXIT_CODE: i32 = 73; +const HEADER_BYTES: usize = 12; +const MAGIC: &[u8; 4] = b"OSUP"; +// JSON can expand one control character into six bytes (a Unicode escape). +const MAX_ERROR_BYTES: usize = MAX_DIAGNOSTIC_BYTES * 6 + 512; + +#[derive(Debug, thiserror::Error)] +pub enum ProtocolError { + #[error("upload transport: {0}")] + Io(#[from] std::io::Error), + #[error("invalid upload protocol: {0}")] + Invalid(&'static str), + #[error("invalid upload metadata: {0}")] + Json(#[from] serde_json::Error), +} + +impl ProtocolError { + #[must_use] + pub const fn is_io(&self) -> bool { + matches!(self, Self::Io(_)) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum DestinationKind { + Directory, + File, + Missing, +} + +impl DestinationKind { + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::Directory => "directory", + Self::File => "file", + Self::Missing => "missing", + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ExpectedDestination { + pub path: String, + pub kind: DestinationKind, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UploadPlan { + pub destination_directory: String, + pub expected_destination: Option, +} + +impl UploadPlan { + pub fn validate(&self) -> Result<(), ProtocolError> { + validate_path(&self.destination_directory)?; + if let Some(expected) = &self.expected_destination { + validate_path(&expected.path)?; + } + Ok(()) + } +} + +fn validate_path(path: &str) -> Result<(), ProtocolError> { + if path.is_empty() || path.len() > MAX_METADATA_BYTES || path.contains('\0') { + return Err(ProtocolError::Invalid("invalid destination path")); + } + Ok(()) +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct UploadSummary { + pub bytes: u64, + pub sha256: [u8; 32], +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum FailureCode { + Protocol, + DestinationChanged, + Extract, + InputWrite, + Internal, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct UploadFailure { + pub code: FailureCode, + pub message: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Frame { + Hello { + window_bytes: u32, + max_data_bytes: u32, + }, + Begin(UploadPlan), + Ready, + Data(Vec), + Ack(u64), + End(UploadSummary), + Result(UploadSummary), + Error(UploadFailure), + Cancel, +} + +impl Frame { + pub fn encode(&self) -> Result, ProtocolError> { + let (kind, payload) = match self { + Self::Hello { + window_bytes, + max_data_bytes, + } => { + validate_limits(*window_bytes, *max_data_bytes)?; + let mut bytes = window_bytes.to_be_bytes().to_vec(); + bytes.extend_from_slice(&max_data_bytes.to_be_bytes()); + (1, bytes) + } + Self::Begin(plan) => { + plan.validate()?; + (2, serde_json::to_vec(plan)?) + } + Self::Ready => (3, Vec::new()), + Self::Data(bytes) => (4, bytes.clone()), + Self::Ack(bytes) => (5, bytes.to_be_bytes().to_vec()), + Self::End(summary) => (6, encode_summary(*summary)), + Self::Result(summary) => (7, encode_summary(*summary)), + Self::Error(failure) => { + if failure.message.len() > MAX_DIAGNOSTIC_BYTES { + return Err(ProtocolError::Invalid("diagnostic exceeds limit")); + } + (8, serde_json::to_vec(failure)?) + } + Self::Cancel => (9, Vec::new()), + }; + validate_length(kind, payload.len())?; + let length = u32::try_from(payload.len()) + .map_err(|_| ProtocolError::Invalid("frame exceeds limit"))?; + let mut encoded = Vec::with_capacity(HEADER_BYTES + payload.len()); + encoded.extend_from_slice(MAGIC); + encoded.extend_from_slice(&[kind, 0, 0, 0]); + encoded.extend_from_slice(&length.to_be_bytes()); + encoded.extend_from_slice(&payload); + Ok(encoded) + } +} + +fn validate_limits(window: u32, chunk: u32) -> Result<(), ProtocolError> { + if window == 0 + || window > WINDOW_BYTES + || chunk == 0 + || chunk as usize > MAX_DATA_BYTES + || chunk > window + { + return Err(ProtocolError::Invalid("invalid credit limits")); + } + Ok(()) +} + +fn encode_summary(summary: UploadSummary) -> Vec { + let mut encoded = summary.bytes.to_be_bytes().to_vec(); + encoded.extend_from_slice(&summary.sha256); + encoded +} + +fn validate_length(kind: u8, length: usize) -> Result<(), ProtocolError> { + let valid = match kind { + 1 | 5 => length == 8, + 2 => length > 0 && length <= MAX_METADATA_BYTES, + 3 | 9 => length == 0, + 4 => length > 0 && length <= MAX_DATA_BYTES, + 6 | 7 => length == 40, + 8 => length > 0 && length <= MAX_ERROR_BYTES, + _ => return Err(ProtocolError::Invalid("unknown frame type")), + }; + if !valid { + return Err(ProtocolError::Invalid("invalid frame length")); + } + Ok(()) +} + +fn parse_header(header: &[u8; HEADER_BYTES]) -> Result<(u8, usize), ProtocolError> { + if &header[..4] != MAGIC || header[5..8] != [0, 0, 0] { + return Err(ProtocolError::Invalid("invalid frame header")); + } + let kind = header[4]; + let length = u32::from_be_bytes(header[8..12].try_into().expect("fixed length")) as usize; + validate_length(kind, length)?; + Ok((kind, length)) +} + +fn decode_payload(kind: u8, payload: Vec) -> Result { + Ok(match kind { + 1 => { + let window_bytes = u32::from_be_bytes(payload[..4].try_into().expect("fixed length")); + let max_data_bytes = + u32::from_be_bytes(payload[4..8].try_into().expect("fixed length")); + validate_limits(window_bytes, max_data_bytes)?; + Frame::Hello { + window_bytes, + max_data_bytes, + } + } + 2 => { + let plan: UploadPlan = serde_json::from_slice(&payload)?; + plan.validate()?; + Frame::Begin(plan) + } + 3 => Frame::Ready, + 4 => Frame::Data(payload), + 5 => Frame::Ack(u64::from_be_bytes( + payload[..8].try_into().expect("fixed length"), + )), + 6 | 7 => { + let summary = UploadSummary { + bytes: u64::from_be_bytes(payload[..8].try_into().expect("fixed length")), + sha256: payload[8..40].try_into().expect("fixed length"), + }; + if kind == 6 { + Frame::End(summary) + } else { + Frame::Result(summary) + } + } + 8 => { + let failure: UploadFailure = serde_json::from_slice(&payload)?; + if failure.message.len() > MAX_DIAGNOSTIC_BYTES { + return Err(ProtocolError::Invalid("diagnostic exceeds limit")); + } + Frame::Error(failure) + } + 9 => Frame::Cancel, + _ => return Err(ProtocolError::Invalid("unknown frame type")), + }) +} + +/// Incremental decoder holding at most one bounded payload. The caller keeps +/// processing the remaining slice, so coalesced packets do not grow a queue. +#[derive(Default)] +pub struct Decoder { + header: [u8; HEADER_BYTES], + header_used: usize, + payload: Vec, + expected: Option<(u8, usize)>, +} + +impl Decoder { + pub fn decode(&mut self, input: &mut &[u8]) -> Result, ProtocolError> { + if self.header_used < HEADER_BYTES { + let size = input.len().min(HEADER_BYTES - self.header_used); + self.header[self.header_used..self.header_used + size].copy_from_slice(&input[..size]); + self.header_used += size; + *input = &input[size..]; + if self.header_used < HEADER_BYTES { + return Ok(None); + } + let expected = parse_header(&self.header).inspect_err(|_| { + self.header_used = 0; + })?; + self.payload.reserve(expected.1); + self.expected = Some(expected); + } + let (kind, length) = self.expected.expect("validated header"); + let size = input.len().min(length - self.payload.len()); + self.payload.extend_from_slice(&input[..size]); + *input = &input[size..]; + if self.payload.len() < length { + return Ok(None); + } + self.header_used = 0; + self.expected = None; + decode_payload(kind, std::mem::take(&mut self.payload)).map(Some) + } + + #[must_use] + pub const fn is_empty(&self) -> bool { + self.header_used == 0 + } +} + +pub async fn read_frame(reader: &mut R) -> Result { + let mut header = [0_u8; HEADER_BYTES]; + reader.read_exact(&mut header).await?; + let (kind, length) = parse_header(&header)?; + let mut payload = vec![0; length]; + reader.read_exact(&mut payload).await?; + decode_payload(kind, payload) +} + +pub async fn write_frame( + writer: &mut W, + frame: &Frame, +) -> Result<(), ProtocolError> { + writer.write_all(&frame.encode()?).await?; + Ok(()) +} + +// This is the existing CLI destination probe, now shared with the receiving +// supervisor. Check existing parents before treating failed existence as absence. +const DESTINATION_KIND_SH: &str = r#"upload_destination_kind() { + upload_path=$1 + upload_parent=$2 + while :; do + if [ -d "$upload_parent" ]; then + if [ ! -x "$upload_parent" ]; then + printf 'cannot inspect upload destination parent: %s\n' "$upload_parent" >&2 + return 1 + fi + elif [ -e "$upload_parent" ] || [ -L "$upload_parent" ]; then + printf 'upload destination parent is not a directory: %s\n' "$upload_parent" >&2 + return 1 + fi + case "$upload_parent" in + /|.) break ;; + */*) upload_parent=${upload_parent%/*}; [ -n "$upload_parent" ] || upload_parent=/ ;; + *) upload_parent=. ;; + esac + done + if [ -d "$upload_path" ]; then + printf directory + elif [ -e "$upload_path" ] || [ -L "$upload_path" ]; then + printf file + else + printf missing + fi +}"#; + +fn destination_invocation(path: &str) -> String { + let parent = match path.rfind('/') { + Some(0) => "/", + Some(position) => &path[..position], + None => ".", + }; + format!( + "upload_destination_kind {} {}", + shell_escape(path), + shell_escape(parent) + ) +} + +#[must_use] +pub fn build_destination_probe(path: &str) -> String { + format!("{DESTINATION_KIND_SH}\n{}", destination_invocation(path)) +} + +pub fn build_extract_command(plan: &UploadPlan) -> Result { + plan.validate()?; + let dest = shell_escape(&plan.destination_directory); + let extract = format!("mkdir -p -- {dest} && cat | tar xf - -C {dest}"); + let Some(expected) = &plan.expected_destination else { + return Ok(extract); + }; + Ok(format!( + "{DESTINATION_KIND_SH}\n\ + upload_kind=$({probe}) || exit 1\n\ + if [ \"$upload_kind\" != {kind} ]; then\n\ + printf 'upload destination changed type before extraction: %s\\n' {path} >&2\n\ + exit {status}\n\ + fi\n\ + {extract}", + probe = destination_invocation(&expected.path), + kind = expected.kind.as_str(), + path = shell_escape(&expected.path), + status = DESTINATION_CHANGED_EXIT_CODE, + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn frames() -> Vec { + let summary = UploadSummary { + bytes: u64::MAX, + sha256: [7; 32], + }; + vec![ + Frame::Hello { + window_bytes: WINDOW_BYTES, + max_data_bytes: u32::try_from(MAX_DATA_BYTES).unwrap(), + }, + Frame::Begin(UploadPlan { + destination_directory: "/tmp/with 'quotes'".into(), + expected_destination: Some(ExpectedDestination { + path: "/tmp/file".into(), + kind: DestinationKind::Missing, + }), + }), + Frame::Ready, + Frame::Data(vec![42; MAX_DATA_BYTES]), + Frame::Ack(u64::MAX), + Frame::End(summary), + Frame::Result(summary), + Frame::Error(UploadFailure { + code: FailureCode::InputWrite, + message: "stdin closed".into(), + }), + Frame::Cancel, + ] + } + + #[test] + fn fragmented_and_coalesced_frames_round_trip() { + let frames = frames(); + let encoded = frames + .iter() + .flat_map(|f| f.encode().unwrap()) + .collect::>(); + for size in [1, 7, 12, 4096, encoded.len()] { + let mut decoder = Decoder::default(); + let mut output = Vec::new(); + for chunk in encoded.chunks(size) { + let mut input = chunk; + while !input.is_empty() { + if let Some(frame) = decoder.decode(&mut input).unwrap() { + output.push(frame); + } + } + } + assert!(decoder.is_empty()); + assert_eq!(output, frames); + } + } + + #[test] + fn malformed_headers_fail_before_payload_allocation() { + let good = Frame::Ready.encode().unwrap(); + for offset in [0, 4, 5, 6, 7, 8, 9, 10, 11] { + let mut bad = good.clone(); + bad[offset] = 255; + let mut decoder = Decoder::default(); + assert!(decoder.decode(&mut bad.as_slice()).is_err()); + assert_eq!(decoder.payload.capacity(), 0); + } + assert!(Frame::Data(vec![0; MAX_DATA_BYTES + 1]).encode().is_err()); + assert!( + Frame::Hello { + window_bytes: WINDOW_BYTES + 1, + max_data_bytes: 1 + } + .encode() + .is_err() + ); + } + + #[test] + fn metadata_rejects_unknown_fields_and_nul_paths() { + let bad = br#"{"destination_directory":"/tmp","expected_destination":null,"command":"id"}"# + .to_vec(); + assert!(decode_payload(2, bad).is_err()); + assert!( + UploadPlan { + destination_directory: "/tmp\0bad".into(), + expected_destination: None + } + .validate() + .is_err() + ); + assert!( + Frame::Error(UploadFailure { + code: FailureCode::Protocol, + message: "a".repeat(MAX_DIAGNOSTIC_BYTES + 1) + }) + .encode() + .is_err() + ); + } + + #[tokio::test] + async fn async_codec_rejects_truncation() { + let encoded = Frame::Data(vec![7; 32]).encode().unwrap(); + let mut input = &encoded[..encoded.len() - 1]; + assert!(matches!( + read_frame(&mut input).await, + Err(ProtocolError::Io(_)) + )); + let mut output = Vec::new(); + write_frame(&mut output, &Frame::Ack(12)).await.unwrap(); + assert_eq!( + read_frame(&mut output.as_slice()).await.unwrap(), + Frame::Ack(12) + ); + } +} diff --git a/crates/openshell-supervisor-process/README.md b/crates/openshell-supervisor-process/README.md new file mode 100644 index 0000000000..7d6e6f13db --- /dev/null +++ b/crates/openshell-supervisor-process/README.md @@ -0,0 +1,31 @@ +# OpenShell process supervisor + +This crate manages workload processes and SSH sessions through the admitted +isolation backend. + +## Bounded streaming uploads + +The CLI's built-in uploads use the `openshell-upload-v1` SSH subsystem. The +gateway relays its SSH channel without interpreting upload frames. The +supervisor launches the extractor through the ordinary `BoundaryExec` path, +with the same workload identity, environment, working directory, and +filesystem policy as other SSH children. + +The upload protocol reserves stdout for control frames and the final result. +Extractor output is drained into bounded diagnostics. The supervisor grants a +1 MiB application window, and data frames carry at most 64 KiB. The CLI stops +sending archive data when outstanding bytes reach the window. A cumulative +acknowledgment returns credit only after `InputTask` successfully writes those +bytes through `BoundaryInput::write_all` and clears them from its in-flight +budget. The existing 4 MiB pending-input guard remains active. + +The end frame carries the archive byte count and SHA-256 digest. A successful +result requires a matching stream, drained accepted writes, and a successful +extractor exit. Protocol failure, extractor failure, or channel cancellation +terminates the upload process and stops granting credit. Extraction writes +directly to the destination, so failures can leave partially extracted files; +the protocol does not provide rollback or atomic replacement. + +The CLI requires supervisor support for this subsystem before sending archive +data. It does not fall back to the previous unpaced exec-stdin upload path. +Ordinary exec stdin and SFTP retain their existing input-budget behavior. diff --git a/crates/openshell-supervisor-process/src/ssh.rs b/crates/openshell-supervisor-process/src/ssh.rs index 75c7b1fa5c..18fd521ab3 100644 --- a/crates/openshell-supervisor-process/src/ssh.rs +++ b/crates/openshell-supervisor-process/src/ssh.rs @@ -4,10 +4,14 @@ //! Embedded SSH server for sandbox access. mod input; +mod upload; #[cfg(test)] mod exec_input_tests; +#[cfg(test)] +mod upload_tests; + use crate::main_session::{MainOutput, MainSession}; #[cfg(unix)] use libc; @@ -371,6 +375,8 @@ async fn handle_connection( #[allow(clippy::struct_excessive_bools)] #[derive(Default)] struct ChannelState { + upload: Option, + command_started: bool, input_sender: Option, input_task: Option, process: Option>, @@ -607,6 +613,10 @@ impl russh::server::Handler for SshHandler { .channels .get_mut(&channel) .ok_or_else(|| anyhow::anyhow!("pty_request on unknown channel {channel:?}"))?; + if state.upload.is_some() || state.command_started { + session.channel_failure(channel)?; + return Ok(()); + } state.pty_request = Some(PtyRequest { term: term.to_string(), col_width, @@ -652,6 +662,14 @@ impl russh::server::Handler for SshHandler { channel: ChannelId, session: &mut Session, ) -> Result<(), Self::Error> { + if self + .channels + .get(&channel) + .is_some_and(|state| state.upload.is_some() || state.command_started) + { + session.channel_failure(channel)?; + return Ok(()); + } // Only allocate a PTY when the client explicitly requested one via // pty_request. VS Code Remote-SSH sends shell_request *without* a // preceding pty_request and expects pipe-based I/O with clean LF line @@ -674,6 +692,14 @@ impl russh::server::Handler for SshHandler { data: &[u8], session: &mut Session, ) -> Result<(), Self::Error> { + if self + .channels + .get(&channel) + .is_some_and(|state| state.upload.is_some() || state.command_started) + { + session.channel_failure(channel)?; + return Ok(()); + } let command = String::from_utf8_lossy(data).trim().to_string(); if command.is_empty() { session.channel_success(channel)?; @@ -698,6 +724,28 @@ impl russh::server::Handler for SshHandler { name: &str, session: &mut Session, ) -> Result<(), Self::Error> { + let Some(state) = self.channels.get_mut(&channel) else { + session.channel_failure(channel)?; + return Ok(()); + }; + if state.upload.is_some() || state.command_started || state.main_attached { + session.channel_failure(channel)?; + return Ok(()); + } + if name == openshell_core::upload::SUBSYSTEM { + if state.pty_request.is_some() { + session.channel_failure(channel)?; + return Ok(()); + } + state.upload = Some(upload::UploadChannel::spawn( + channel, + session.handle(), + self.boundary_exec.clone(), + !state.no_login_shell, + )); + session.channel_success(channel)?; + return Ok(()); + } if name == "openshell-main" { let Some(state) = self.channels.get(&channel) else { session.channel_failure(channel)?; @@ -881,6 +929,14 @@ impl russh::server::Handler for SshHandler { data: &[u8], session: &mut Session, ) -> Result<(), Self::Error> { + if let Some(upload) = self + .channels + .get_mut(&channel) + .and_then(|state| state.upload.as_mut()) + { + upload.data(data); + return Ok(()); + } let Some(state) = self.channels.get_mut(&channel) else { warn!("data on unknown channel {channel:?}"); return Ok(()); @@ -951,6 +1007,10 @@ impl russh::server::Handler for SshHandler { // closes stdin. Commands such as `cat | tar xf -` need this half-close // to finish while their output remains available to the SSH client. if let Some(state) = self.channels.get_mut(&channel) { + if let Some(upload) = state.upload.as_mut() { + upload.eof(); + return Ok(()); + } if state.main_attached && let Some(owner) = state.main_input_owner.take() && let Some(main_session) = self.main_session.as_ref() @@ -969,6 +1029,23 @@ impl russh::server::Handler for SshHandler { Ok(()) } + async fn extended_data( + &mut self, + channel: ChannelId, + _ext: u32, + _data: &[u8], + _session: &mut Session, + ) -> Result<(), Self::Error> { + if let Some(upload) = self + .channels + .get_mut(&channel) + .and_then(|state| state.upload.as_mut()) + { + upload.reject_extended_data(); + } + Ok(()) + } + async fn signal( &mut self, channel: ChannelId, @@ -995,11 +1072,12 @@ impl russh::server::Handler for SshHandler { } return Ok(()); } - let Some(process) = self - .channels - .get(&channel) - .and_then(|state| state.process.clone()) - else { + let Some(process) = self.channels.get(&channel).and_then(|state| { + state + .process + .clone() + .or_else(|| state.upload.as_ref()?.process()) + }) else { return Ok(()); }; let signal = match signal { @@ -1112,6 +1190,7 @@ impl SshHandler { .get_mut(&channel) .ok_or_else(|| anyhow::anyhow!("exec on unknown channel {channel:?}"))?; state.process = Some(exec.process.clone()); + state.command_started = true; state.terminal = exec.terminal.take(); let output_status = exec.output_status.take(); diff --git a/crates/openshell-supervisor-process/src/ssh/input.rs b/crates/openshell-supervisor-process/src/ssh/input.rs index 391fa62b70..46ca1dfa85 100644 --- a/crates/openshell-supervisor-process/src/ssh/input.rs +++ b/crates/openshell-supervisor-process/src/ssh/input.rs @@ -8,7 +8,7 @@ use std::sync::{Arc, Mutex}; use openshell_isolation_interface::contract::BoundaryInput; use tokio::io::AsyncWriteExt; -use tokio::sync::Notify; +use tokio::sync::{Notify, watch}; use tokio::task::{AbortHandle, JoinHandle}; /// Includes queued bytes and the chunk currently being written. Consumed bytes @@ -28,6 +28,22 @@ struct Buffer { struct Shared { buffer: Mutex, ready: Notify, + progress: Option>, +} + +/// A watch keeps only the newest consumed offset even when SSH output stalls. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(super) struct InputProgress { + pub(super) consumed: u64, + pub(super) state: InputState, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(super) enum InputState { + #[default] + Writing, + Drained, + Failed, } pub(super) struct InputSender(Arc); @@ -73,8 +89,30 @@ pub(super) struct InputTask { impl InputTask { pub(super) fn spawn(stdin: BoundaryInput) -> (InputSender, Self) { - let shared = Arc::new(Shared::default()); - let receiver = Receiver(shared.clone()); + Self::spawn_inner(stdin, None) + } + + pub(super) fn spawn_with_progress( + stdin: BoundaryInput, + ) -> (InputSender, Self, watch::Receiver) { + let (progress, receiver) = watch::channel(InputProgress::default()); + let (sender, task) = Self::spawn_inner(stdin, Some(progress)); + (sender, task, receiver) + } + + fn spawn_inner( + stdin: BoundaryInput, + progress: Option>, + ) -> (InputSender, Self) { + let shared = Arc::new(Shared { + progress, + ..Shared::default() + }); + let receiver = Receiver { + shared: shared.clone(), + consumed: 0, + finished: false, + }; let task = tokio::spawn(receiver.write_to(stdin)); (InputSender(shared.clone()), Self { shared, task }) } @@ -96,15 +134,28 @@ impl Drop for InputTask { } } -struct Receiver(Arc); +struct Receiver { + shared: Arc, + consumed: u64, + finished: bool, +} impl Receiver { - async fn write_to(self, mut stdin: BoundaryInput) { + fn publish(&self, state: InputState) { + if let Some(progress) = &self.shared.progress { + progress.send_replace(InputProgress { + consumed: self.consumed, + state, + }); + } + } + + async fn write_to(mut self, mut stdin: BoundaryInput) { loop { - let ready = self.0.ready.notified(); + let ready = self.shared.ready.notified(); let chunk = { - let Ok(mut buffer) = self.0.buffer.lock() else { - break; + let Ok(mut buffer) = self.shared.buffer.lock() else { + return; }; if buffer.bytes.is_empty() && buffer.closed { break; @@ -118,12 +169,22 @@ impl Receiver { continue; } if stdin.write_all(&chunk).await.is_err() { - break; + return; } - let Ok(mut buffer) = self.0.buffer.lock() else { - break; - }; - buffer.in_flight = 0; + { + let Ok(mut buffer) = self.shared.buffer.lock() else { + return; + }; + buffer.in_flight = 0; + } + self.consumed += chunk.len() as u64; + self.publish(InputState::Writing); + } + // Ordinary exec retains its existing drop-to-close behavior. Upload + // needs an explicit shutdown result before claiming the drain complete. + if self.shared.progress.is_none() || stdin.shutdown().await.is_ok() { + self.finished = true; + self.publish(InputState::Drained); } } } @@ -132,11 +193,14 @@ impl Drop for Receiver { fn drop(&mut self) { // A failed or cancelled writer must refuse later input and release the // queue even if the SSH channel has not yet received a close packet. - if let Ok(mut buffer) = self.0.buffer.lock() { + if let Ok(mut buffer) = self.shared.buffer.lock() { buffer.closed = true; buffer.in_flight = 0; buffer.bytes = VecDeque::new(); } + if !self.finished { + self.publish(InputState::Failed); + } } } @@ -182,4 +246,79 @@ mod tests { .unwrap(); assert_eq!(received, vec![42; 8192]); } + + #[tokio::test] + async fn progress_is_earned_after_complete_writes_and_reports_drain() { + let (stdin, mut reader) = tokio::io::duplex(1); + let (sender, _task, mut progress) = InputTask::spawn_with_progress(Box::new(stdin)); + sender.send(&vec![42; WRITE_CHUNK]).unwrap(); + reader.read_u8().await.unwrap(); + assert_eq!(progress.borrow().consumed, 0); + drop(sender); + let mut remaining = Vec::new(); + reader.read_to_end(&mut remaining).await.unwrap(); + assert_eq!(remaining, vec![42; WRITE_CHUNK - 1]); + progress + .wait_for(|progress| progress.state == InputState::Drained) + .await + .unwrap(); + assert_eq!(progress.borrow().consumed, WRITE_CHUNK as u64); + } + + #[tokio::test] + async fn failed_writer_publishes_terminal_state_and_refuses_input() { + let (stdin, reader) = tokio::io::duplex(1); + let (sender, _task, mut progress) = InputTask::spawn_with_progress(Box::new(stdin)); + drop(reader); + sender.send(b"write fails").unwrap(); + progress + .wait_for(|progress| progress.state == InputState::Failed) + .await + .unwrap(); + assert_eq!(progress.borrow().consumed, 0); + assert_eq!(sender.send(b"later input"), Err(SendError::Closed)); + } + + #[tokio::test] + async fn shutdown_failure_preserves_consumption_without_claiming_drain() { + struct ShutdownFails(tokio::io::DuplexStream); + + impl tokio::io::AsyncWrite for ShutdownFails { + fn poll_write( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + bytes: &[u8], + ) -> std::task::Poll> { + std::pin::Pin::new(&mut self.0).poll_write(cx, bytes) + } + + fn poll_flush( + mut self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::pin::Pin::new(&mut self.0).poll_flush(cx) + } + + fn poll_shutdown( + self: std::pin::Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Err(std::io::Error::other("shutdown failed"))) + } + } + + let (stdin, mut reader) = tokio::io::duplex(1024); + let (sender, _task, mut progress) = + InputTask::spawn_with_progress(Box::new(ShutdownFails(stdin))); + sender.send(b"consumed before shutdown").unwrap(); + drop(sender); + let mut received = Vec::new(); + reader.read_to_end(&mut received).await.unwrap(); + assert_eq!(received, b"consumed before shutdown"); + progress + .wait_for(|progress| progress.state == InputState::Failed) + .await + .unwrap(); + assert_eq!(progress.borrow().consumed, received.len() as u64); + } } diff --git a/crates/openshell-supervisor-process/src/ssh/upload.rs b/crates/openshell-supervisor-process/src/ssh/upload.rs new file mode 100644 index 0000000000..cb0453051d --- /dev/null +++ b/crates/openshell-supervisor-process/src/ssh/upload.rs @@ -0,0 +1,586 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Upload input credit follows boundary consumption, independently of russh's +//! receive window. The SSH callback only decodes and queues bounded input. + +use super::input::{InputProgress, InputSender, InputState, InputTask, SendError}; +use openshell_core::upload::{ + DESTINATION_CHANGED_EXIT_CODE, Decoder, FailureCode, Frame, MAX_DATA_BYTES, + MAX_DIAGNOSTIC_BYTES, UploadFailure, UploadPlan, UploadSummary, WINDOW_BYTES, + build_extract_command, +}; +use openshell_isolation_interface::contract::{ + BoundaryExec, BoundaryExitStatus, BoundaryProcess, ExecSession, ExecSpec, ShellSpec, +}; +use russh::ChannelId; +use russh::server::Handle; +use sha2::{Digest, Sha256}; +use std::sync::{Arc, Mutex}; +use tokio::io::{AsyncRead, AsyncReadExt}; +use tokio::sync::{Notify, oneshot, watch}; +use tokio::task::JoinHandle; + +const ACK_INTERVAL: u64 = 256 * 1024; + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Phase { + AwaitBegin, + Starting, + Receiving, + Draining, + Finished, +} + +struct State { + phase: Phase, + plan: Option, + sender: Option, + input_task: Option, + progress: Option>, + process: Option>, + received: u64, + digest: Sha256, + summary: Option, + failure: Option, +} + +struct Shared { + state: Mutex, + changed: Notify, +} + +pub(super) struct UploadChannel { + decoder: Decoder, + shared: Arc, + owner: JoinHandle<()>, +} + +impl UploadChannel { + pub(super) fn spawn( + channel: ChannelId, + handle: Handle, + boundary_exec: Arc, + login: bool, + ) -> Self { + let shared = Arc::new(Shared { + state: Mutex::new(State { + phase: Phase::AwaitBegin, + plan: None, + sender: None, + input_task: None, + progress: None, + process: None, + received: 0, + digest: Sha256::new(), + summary: None, + failure: None, + }), + changed: Notify::new(), + }); + let owner_shared = shared.clone(); + let owner = tokio::spawn(async move { + run(owner_shared, channel, handle, boundary_exec, login).await; + }); + Self { + decoder: Decoder::default(), + shared, + owner, + } + } + + pub(super) fn data(&mut self, mut bytes: &[u8]) { + while !bytes.is_empty() { + match self.decoder.decode(&mut bytes) { + Ok(Some(frame)) => { + if let Err(failure) = self.accept(frame) { + self.fail(failure); + return; + } + } + Ok(None) => return, + Err(error) => { + self.fail(failure(FailureCode::Protocol, error.to_string())); + return; + } + } + } + } + + pub(super) fn eof(&self) { + let valid = self.decoder.is_empty() + && self + .shared + .state + .lock() + .is_ok_and(|state| matches!(state.phase, Phase::Draining | Phase::Finished)); + if !valid { + self.fail(failure( + FailureCode::Protocol, + "upload stream ended before a complete END frame", + )); + } + } + + pub(super) fn reject_extended_data(&self) { + self.fail(failure( + FailureCode::Protocol, + "extended data is not supported by the upload subsystem", + )); + } + + pub(super) fn process(&self) -> Option> { + self.shared.state.lock().ok()?.process.clone() + } + + fn fail(&self, failure: UploadFailure) { + if let Ok(mut state) = self.shared.state.lock() + && state.phase != Phase::Finished + && state.failure.is_none() + { + state.failure = Some(failure); + // Stop a stalled write immediately. The owner reports the typed + // failure and terminates the process outside the session callback. + state.sender.take(); + state.input_task.take(); + } + self.shared.changed.notify_one(); + } + + fn accept(&self, frame: Frame) -> Result<(), UploadFailure> { + let mut state = self + .shared + .state + .lock() + .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; + if state.failure.is_some() || state.phase == Phase::Finished { + return Ok(()); + } + match frame { + Frame::Begin(plan) if state.phase == Phase::AwaitBegin => { + plan.validate() + .map_err(|error| failure(FailureCode::Protocol, error.to_string()))?; + state.phase = Phase::Starting; + state.plan = Some(plan); + } + Frame::Data(data) if state.phase == Phase::Receiving => { + let consumed = state + .progress + .as_ref() + .map_or(0, |progress| progress.borrow().consumed); + let received = state + .received + .checked_add(data.len() as u64) + .ok_or_else(|| failure(FailureCode::Protocol, "upload length overflow"))?; + if received - consumed > u64::from(WINDOW_BYTES) { + return Err(failure( + FailureCode::Protocol, + "upload input credit exceeded", + )); + } + let sender = state.sender.as_ref().ok_or_else(|| { + failure( + FailureCode::InputWrite, + "extraction process stdin is closed", + ) + })?; + sender.send(&data).map_err(|error| match error { + SendError::Full => { + failure(FailureCode::Internal, "upload input budget exceeded") + } + SendError::Closed => failure( + FailureCode::InputWrite, + "extraction process stdin is closed", + ), + })?; + state.received = received; + state.digest.update(data); + } + Frame::End(summary) if state.phase == Phase::Receiving => { + let actual = UploadSummary { + bytes: state.received, + sha256: state.digest.clone().finalize().into(), + }; + if summary != actual { + return Err(failure( + FailureCode::Protocol, + "upload count or digest mismatch", + )); + } + state.summary = Some(summary); + state.phase = Phase::Draining; + // EOF drains accepted bytes before closing boundary stdin. + state.sender.take(); + } + Frame::Cancel => { + return Err(failure(FailureCode::Protocol, "upload cancelled by client")); + } + _ => { + return Err(failure(FailureCode::Protocol, "unexpected upload frame")); + } + } + drop(state); + self.shared.changed.notify_one(); + Ok(()) + } +} + +impl Drop for UploadChannel { + fn drop(&mut self) { + self.owner.abort(); + let process = self.shared.state.lock().ok().and_then(|mut state| { + state.phase = Phase::Finished; + state.sender.take(); + state.input_task.take(); + state.process.take() + }); + if let Some(process) = process { + tokio::spawn(async move { + let _ = process.terminate().await; + }); + } + } +} + +fn failure(code: FailureCode, message: impl Into) -> UploadFailure { + let mut message = message.into(); + // Error formatting and backend diagnostics can be large. Bound the UTF-8 + // string before it enters a queued protocol frame. + if message.len() > MAX_DIAGNOSTIC_BYTES { + let mut end = MAX_DIAGNOSTIC_BYTES; + while !message.is_char_boundary(end) { + end -= 1; + } + message.truncate(end); + } + UploadFailure { code, message } +} + +async fn send(handle: &Handle, channel: ChannelId, frame: Frame) -> bool { + match frame.encode() { + Ok(bytes) => handle.data(channel, bytes).await.is_ok(), + Err(_) => false, + } +} + +async fn finish( + shared: &Shared, + channel: ChannelId, + handle: &Handle, + result: Result, +) { + let failed = result.is_err(); + let process = shared.state.lock().ok().and_then(|mut state| { + state.phase = Phase::Finished; + state.sender.take(); + state.input_task.take(); + state.process.take() + }); + if failed && let Some(process) = process { + tokio::spawn(async move { + let _ = process.terminate().await; + }); + } + let frame = match result { + Ok(summary) => Frame::Result(summary), + Err(error) => Frame::Error(error), + }; + let _ = send(handle, channel, frame).await; + let _ = handle + .exit_status_request(channel, if failed { 74 } else { 0 }) + .await; + let _ = handle.eof(channel).await; + let _ = handle.close(channel).await; +} + +async fn run( + shared: Arc, + channel: ChannelId, + handle: Handle, + boundary_exec: Arc, + login: bool, +) { + let result = drive(&shared, channel, &handle, boundary_exec, login).await; + finish(&shared, channel, &handle, result).await; +} + +async fn drive( + shared: &Shared, + channel: ChannelId, + handle: &Handle, + boundary_exec: Arc, + login: bool, +) -> Result { + if !send( + handle, + channel, + Frame::Hello { + window_bytes: WINDOW_BYTES, + max_data_bytes: u32::try_from(MAX_DATA_BYTES).expect("bounded data frame"), + }, + ) + .await + { + return Err(failure( + FailureCode::Internal, + "upload reply delivery failed", + )); + } + let plan = loop { + let changed = shared.changed.notified(); + { + let mut state = shared + .state + .lock() + .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; + if let Some(error) = state.failure.take() { + return Err(error); + } + if let Some(plan) = state.plan.take() { + break plan; + } + } + changed.await; + }; + let command = build_extract_command(&plan) + .map_err(|error| failure(FailureCode::Protocol, error.to_string()))?; + // Remote backends may admit a process before returning its attachment. + // Keep startup owned independently of this channel: cancellation must + // terminate even a result delivered after the upload owner is dropped. + let (started_tx, started_rx) = oneshot::channel(); + tokio::spawn(async move { + let result = boundary_exec + .exec(ExecSpec { + program: String::new(), + args: Vec::new(), + shell: Some(ShellSpec { + command: Some(command), + login, + }), + runtime_helper: None, + env: Vec::new(), + workdir: None, + pty: false, + }) + .await + .map(|exec| PendingExec(Some(exec))); + // A queued PendingExec also terminates on drop if the receiver was + // cancelled after send succeeded but before receiving the attachment. + let _ = started_tx.send(result); + }); + let mut pending = started_rx + .await + .map_err(|_| failure(FailureCode::Internal, "extraction startup stopped"))? + .map_err(|error| failure(FailureCode::Extract, error.to_string()))?; + { + let mut state = shared + .state + .lock() + .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; + // Channel drop can run on another worker after the receive completes. + // Keep the startup guard armed until ownership is registered under the + // same lock that cancellation uses to mark the channel finished. + if state.phase == Phase::Finished || state.failure.is_some() { + return Err(failure( + FailureCode::Protocol, + "upload cancelled during startup", + )); + } + state.process = Some( + pending + .0 + .as_ref() + .expect("owned startup result") + .process + .clone(), + ); + } + let mut exec = pending.0.take().expect("registered startup result"); + let stdin = exec + .stdin + .take() + .ok_or_else(|| failure(FailureCode::Internal, "extraction process has no stdin"))?; + let (sender, input_task, mut progress) = InputTask::spawn_with_progress(stdin); + { + let mut state = shared + .state + .lock() + .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; + if state.phase == Phase::Finished { + return Err(failure( + FailureCode::Protocol, + "upload cancelled during startup", + )); + } + if let Some(error) = state.failure.take() { + return Err(error); + } + state.sender = Some(sender); + state.input_task = Some(input_task); + state.progress = Some(progress.clone()); + state.phase = Phase::Receiving; + } + let diagnostic = Arc::new(Mutex::new(Vec::new())); + let stdout = DrainTask::spawn(exec.stdout, diagnostic.clone()); + let stderr = exec + .stderr + .map(|stderr| DrainTask::spawn(stderr, diagnostic.clone())); + let mut exit = Box::pin(exec.process.wait()); + if !send(handle, channel, Frame::Ready).await { + return Err(failure( + FailureCode::Internal, + "upload reply delivery failed", + )); + } + let mut acknowledged = 0; + let status = loop { + let changed = shared.changed.notified(); + { + let mut state = shared + .state + .lock() + .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; + if let Some(error) = state.failure.take() { + return Err(error); + } + } + let latest = *progress.borrow_and_update(); + if latest.state == InputState::Failed { + return Err(failure( + FailureCode::InputWrite, + "failed to write extraction stdin", + )); + } + if latest.consumed - acknowledged >= ACK_INTERVAL + || (latest.state == InputState::Drained && latest.consumed != acknowledged) + { + if !send(handle, channel, Frame::Ack(latest.consumed)).await { + return Err(failure( + FailureCode::Internal, + "upload reply delivery failed", + )); + } + acknowledged = latest.consumed; + } + tokio::select! { + status = &mut exit => { + let status = status.map_err(|error| failure(FailureCode::Extract, error.to_string()))?; + let final_progress = *progress.borrow(); + let (final_summary, final_received) = shared.state.lock().map(|state| (state.summary, state.received)).map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; + if final_summary.is_none() || final_progress.consumed != final_received { + return Err(extract_failure(status, "extraction process exited before upload completed")); + } + // A pipe shutdown can wake the child just before the writer + // publishes Drained. Await that final watch update rather than + // treating the successful EOF race as an early extraction exit. + while progress.borrow().state == InputState::Writing { + progress.changed().await.map_err(|_| failure(FailureCode::InputWrite, "extraction stdin writer stopped"))?; + } + if progress.borrow().state != InputState::Drained { + return Err(failure(FailureCode::InputWrite, "failed to close extraction stdin")); + } + if final_received != acknowledged && !send(handle, channel, Frame::Ack(final_received)).await { + return Err(failure(FailureCode::Internal, "upload reply delivery failed")); + } + break status; + } + () = changed => {} + result = progress.changed() => { + if result.is_err() { + return Err(failure(FailureCode::InputWrite, "extraction stdin writer stopped")); + } + } + } + }; + stdout + .join() + .await + .map_err(|error| failure(FailureCode::Extract, error.to_string()))?; + if let Some(stderr) = stderr { + stderr + .join() + .await + .map_err(|error| failure(FailureCode::Extract, error.to_string()))?; + } + let status = if let Some(output_status) = exec.output_status { + output_status + .await + .map_err(|_| failure(FailureCode::Extract, "extraction output delivery failed"))? + } else { + status + }; + if status != BoundaryExitStatus::Exited(0) { + let diagnostic = diagnostic.lock().map_or_else( + |_| String::new(), + |bytes| String::from_utf8_lossy(&bytes).into_owned(), + ); + return Err(extract_failure(status, &diagnostic)); + } + shared + .state + .lock() + .ok() + .and_then(|state| state.summary) + .ok_or_else(|| failure(FailureCode::Internal, "upload summary unavailable")) +} + +fn extract_failure(status: BoundaryExitStatus, diagnostic: &str) -> UploadFailure { + if status == BoundaryExitStatus::Exited(DESTINATION_CHANGED_EXIT_CODE) { + failure(FailureCode::DestinationChanged, diagnostic) + } else { + failure( + FailureCode::Extract, + format!("extraction failed ({status:?}): {diagnostic}"), + ) + } +} + +/// Own a late startup result until the channel registers process ownership. +struct PendingExec(Option); + +impl Drop for PendingExec { + fn drop(&mut self) { + if let Some(exec) = self.0.take() { + tokio::spawn(async move { + let _ = exec.process.terminate().await; + }); + } + } +} + +/// Aborting the owner also aborts output drains; dropping `JoinHandle` alone +/// would detach a pipe reader and retain the boundary streams after close. +struct DrainTask(JoinHandle>); + +impl DrainTask { + fn spawn( + mut stream: Box, + diagnostic: Arc>>, + ) -> Self { + Self(tokio::spawn(async move { + let mut bytes = [0; 4096]; + loop { + match stream.read(&mut bytes).await { + Ok(0) => return Ok(()), + Err(error) => return Err(error), + Ok(size) => { + if let Ok(mut diagnostic) = diagnostic.lock() { + let retain = size.min(MAX_DIAGNOSTIC_BYTES - diagnostic.len()); + diagnostic.extend_from_slice(&bytes[..retain]); + } + } + } + } + })) + } + + async fn join(mut self) -> std::io::Result<()> { + (&mut self.0).await.map_err(std::io::Error::other)? + } +} + +impl Drop for DrainTask { + fn drop(&mut self) { + self.0.abort(); + } +} diff --git a/crates/openshell-supervisor-process/src/ssh/upload_tests.rs b/crates/openshell-supervisor-process/src/ssh/upload_tests.rs new file mode 100644 index 0000000000..31e75b8e2c --- /dev/null +++ b/crates/openshell-supervisor-process/src/ssh/upload_tests.rs @@ -0,0 +1,574 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Exercise the production upload handler over a real SSH session. + +use super::tests::{AcceptAnyServerKey, test_client}; +use openshell_core::upload::{ + Decoder, FailureCode, Frame, MAX_DATA_BYTES, SUBSYSTEM, UploadPlan, UploadSummary, WINDOW_BYTES, +}; +use openshell_isolation_interface::contract::{ + BackendError, BoundaryExec, BoundaryExitStatus, BoundaryProcess, BoundarySignal, ExecSession, + ExecSpec, +}; +use russh::ChannelMsg; +use sha2::{Digest, Sha256}; +use std::collections::VecDeque; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::io::{AsyncReadExt, AsyncWriteExt, DuplexStream}; +use tokio::sync::{mpsc, oneshot, watch}; + +const DEADLINE: Duration = Duration::from_secs(10); + +struct Process { + status: watch::Sender>, + terminated: watch::Sender, + signals: mpsc::UnboundedSender, +} + +#[async_trait::async_trait] +impl BoundaryProcess for Process { + async fn wait(&self) -> Result { + let mut status = self.status.subscribe(); + Ok(status + .wait_for(Option::is_some) + .await + .unwrap() + .as_ref() + .copied() + .unwrap()) + } + + async fn signal(&self, signal: BoundarySignal) -> Result<(), BackendError> { + let _ = self.signals.send(signal); + Ok(()) + } + + async fn terminate(&self) -> Result<(), BackendError> { + self.terminated.send_replace(true); + self.status + .send_replace(Some(BoundaryExitStatus::Exited(0))); + Ok(()) + } +} + +struct Exec(Mutex>); + +#[async_trait::async_trait] +impl BoundaryExec for Exec { + async fn exec(&self, spec: ExecSpec) -> Result { + assert!(!spec.pty); + assert!(spec.runtime_helper.is_none()); + if spec.shell.as_ref().unwrap().command.as_deref() != Some("cat") { + assert!(spec.shell.as_ref().unwrap().login); + } + Ok(self.0.lock().unwrap().pop_front().unwrap()) + } +} + +struct Control { + stdin: DuplexStream, + stdout: DuplexStream, + stderr: DuplexStream, + process: Arc, + signals: mpsc::UnboundedReceiver, +} + +fn exec_session() -> (ExecSession, Control) { + let (stdin, reader) = tokio::io::duplex(1024); + let (stdout, writer) = tokio::io::duplex(4096); + let (stderr, errors) = tokio::io::duplex(4096); + let (signal_tx, signals) = mpsc::unbounded_channel(); + let process = Arc::new(Process { + status: watch::channel(None).0, + terminated: watch::channel(false).0, + signals: signal_tx, + }); + ( + ExecSession { + process: process.clone(), + stdin: Some(Box::new(stdin)), + stdout: Box::new(stdout), + stderr: Some(Box::new(stderr)), + terminal: None, + output_status: None, + }, + Control { + stdin: reader, + stdout: writer, + stderr: errors, + process, + signals, + }, + ) +} + +type Channel = russh::Channel; + +struct Reader { + decoder: Decoder, + frames: VecDeque, + statuses: Vec, +} + +impl Reader { + fn new() -> Self { + Self { + decoder: Decoder::default(), + frames: VecDeque::new(), + statuses: Vec::new(), + } + } + + async fn frame(&mut self, channel: &mut Channel) -> Frame { + loop { + if let Some(frame) = self.frames.pop_front() { + return frame; + } + match channel.wait().await.expect("channel closed before frame") { + ChannelMsg::Data { data } => { + let mut bytes = data.as_ref(); + while !bytes.is_empty() { + if let Some(frame) = self.decoder.decode(&mut bytes).unwrap() { + self.frames.push_back(frame); + } + } + } + ChannelMsg::ExitStatus { exit_status } => self.statuses.push(exit_status), + ChannelMsg::WindowAdjusted { .. } | ChannelMsg::Success | ChannelMsg::Eof => {} + message => panic!("unexpected SSH message: {message:?}"), + } + } + } +} + +async fn send(channel: &Channel, frame: Frame) { + channel + .data(frame.encode().unwrap().as_slice()) + .await + .unwrap(); +} + +async fn start(client: &russh::client::Handle) -> (Channel, Reader) { + let mut channel = client.channel_open_session().await.unwrap(); + channel.request_subsystem(true, SUBSYSTEM).await.unwrap(); + assert!(matches!(channel.wait().await, Some(ChannelMsg::Success))); + let mut reader = Reader::new(); + assert_eq!( + reader.frame(&mut channel).await, + Frame::Hello { + window_bytes: WINDOW_BYTES, + max_data_bytes: u32::try_from(MAX_DATA_BYTES).unwrap(), + } + ); + send( + &channel, + Frame::Begin(UploadPlan { + destination_directory: "/tmp/upload".to_string(), + expected_destination: None, + }), + ) + .await; + assert_eq!(reader.frame(&mut channel).await, Frame::Ready); + (channel, reader) +} + +#[tokio::test] +async fn stalled_upload_preserves_controls_then_resumes_beyond_four_mib() { + tokio::time::timeout(DEADLINE, async { + let (upload, mut control) = exec_session(); + let (sibling, mut other) = exec_session(); + let client = test_client( + None, + Arc::new(Exec(Mutex::new([upload, sibling].into()))), + russh::client::Config::default(), + ) + .await; + let (mut channel, mut reader) = start(&client).await; + let payload = vec![42; MAX_DATA_BYTES]; + let mut sent = 0_u64; + while sent < u64::from(WINDOW_BYTES) { + send(&channel, Frame::Data(payload.clone())).await; + sent += payload.len() as u64; + } + channel.signal(russh::Sig::INT).await.unwrap(); + assert_eq!(control.signals.recv().await, Some(BoundarySignal::Int)); + assert!( + tokio::time::timeout(Duration::from_millis(50), reader.frame(&mut channel)) + .await + .is_err() + ); + // Stalled upload input must leave sibling channel requests responsive. + let mut sibling = client.channel_open_session().await.unwrap(); + sibling.exec(true, "cat").await.unwrap(); + assert!(matches!(sibling.wait().await, Some(ChannelMsg::Success))); + sibling.data(&b"sibling"[..]).await.unwrap(); + sibling.eof().await.unwrap(); + let mut sibling_input = Vec::new(); + other.stdin.read_to_end(&mut sibling_input).await.unwrap(); + assert_eq!(sibling_input, b"sibling"); + sibling.close().await.unwrap(); + // Output is drained into a bounded diagnostic, never exposed as forged + // protocol stdout or allowed to block extraction on a full pipe. + let forged = Frame::Ack(u64::MAX).encode().unwrap(); + control.stdout.write_all(&forged).await.unwrap(); + let noisy = tokio::spawn(async move { + control + .stderr + .write_all(&vec![9; 1024 * 1024]) + .await + .unwrap(); + drop(control.stderr); + }); + let consumed = tokio::spawn(async move { + let mut bytes = Vec::new(); + control.stdin.read_to_end(&mut bytes).await.unwrap(); + assert_eq!(bytes, vec![42; 5 * 1024 * 1024]); + drop(control.stdout); + control + .process + .status + .send_replace(Some(BoundaryExitStatus::Exited(0))); + }); + let mut ack = 0; + while sent < 5 * 1024 * 1024 { + if sent + MAX_DATA_BYTES as u64 > ack + u64::from(WINDOW_BYTES) { + match reader.frame(&mut channel).await { + Frame::Ack(consumed) => { + assert!(consumed >= ack && consumed <= sent); + ack = consumed; + } + frame => panic!("unexpected transfer reply: {frame:?}"), + } + } else { + send(&channel, Frame::Data(payload.clone())).await; + sent += payload.len() as u64; + } + } + let summary = UploadSummary { + bytes: sent, + sha256: Sha256::digest(vec![42; usize::try_from(sent).unwrap()]).into(), + }; + send(&channel, Frame::End(summary)).await; + channel.eof().await.unwrap(); + loop { + match reader.frame(&mut channel).await { + Frame::Ack(consumed) => assert!(consumed <= sent), + Frame::Result(actual) => { + assert_eq!(actual, summary); + break; + } + frame => panic!("unexpected terminal reply: {frame:?}"), + } + } + let mut statuses = Vec::new(); + while let Some(message) = channel.wait().await { + match message { + ChannelMsg::ExitStatus { exit_status } => statuses.push(exit_status), + ChannelMsg::Close => break, + _ => {} + } + } + assert_eq!(statuses, [0]); + consumed.await.unwrap(); + noisy.await.unwrap(); + }) + .await + .expect("upload did not resume or blocked SSH control traffic"); +} + +#[tokio::test] +async fn credit_violation_and_truncated_eof_are_typed_and_terminate() { + for truncated in [false, true] { + tokio::time::timeout(DEADLINE, async { + let (session, control) = exec_session(); + let client = test_client( + None, + Arc::new(Exec(Mutex::new([session].into()))), + russh::client::Config::default(), + ) + .await; + let (mut channel, mut reader) = start(&client).await; + if truncated { + channel.eof().await.unwrap(); + } else { + for _ in 0..=WINDOW_BYTES as usize / MAX_DATA_BYTES { + send(&channel, Frame::Data(vec![7; MAX_DATA_BYTES])).await; + } + } + assert!(matches!(reader.frame(&mut channel).await, Frame::Error(error) if error.code == FailureCode::Protocol)); + control.process.terminated.subscribe().wait_for(|terminated| *terminated).await.unwrap(); + let mut eof = Vec::new(); + let _ = control.stdin.take(2).read_to_end(&mut eof).await; + }) + .await + .expect("protocol failure retained a stalled input or process"); + } +} + +#[tokio::test] +async fn broken_stdin_and_early_exit_wake_credit_waiter() { + for early_exit in [false, true] { + tokio::time::timeout(DEADLINE, async { + let (session, control) = exec_session(); + let client = test_client( + None, + Arc::new(Exec(Mutex::new([session].into()))), + russh::client::Config::default(), + ) + .await; + let (mut channel, mut reader) = start(&client).await; + if early_exit { + control.process.status.send_replace(Some(BoundaryExitStatus::Exited(1))); + } else { + drop(control.stdin); + send(&channel, Frame::Data(vec![7; MAX_DATA_BYTES])).await; + } + let expected = if early_exit { FailureCode::Extract } else { FailureCode::InputWrite }; + assert!(matches!(reader.frame(&mut channel).await, Frame::Error(error) if error.code == expected)); + control.process.terminated.subscribe().wait_for(|terminated| *terminated).await.unwrap(); + }) + .await + .expect("terminal failure did not wake upload sender"); + } +} + +#[tokio::test] +async fn close_and_disconnect_cancel_stalled_boundary_input() { + for disconnect in [false, true] { + tokio::time::timeout(DEADLINE, async { + let (session, mut control) = exec_session(); + let client = test_client( + None, + Arc::new(Exec(Mutex::new([session].into()))), + russh::client::Config::default(), + ) + .await; + let (channel, _) = start(&client).await; + send(&channel, Frame::Data(vec![7; MAX_DATA_BYTES])).await; + channel.signal(russh::Sig::INT).await.unwrap(); + control.signals.recv().await.unwrap(); + if disconnect { + client + .disconnect(russh::Disconnect::ByApplication, "test", "") + .await + .unwrap(); + } else { + channel.close().await.unwrap(); + } + control + .process + .terminated + .subscribe() + .wait_for(|terminated| *terminated) + .await + .unwrap(); + let mut received = Vec::new(); + control.stdin.read_to_end(&mut received).await.unwrap(); + assert!(received.len() <= 1024); + }) + .await + .expect("close or disconnect stranded upload input"); + } +} + +#[tokio::test] +async fn upload_rejects_pty_and_duplicate_channel_requests() { + tokio::time::timeout(DEADLINE, async { + let client = test_client( + None, + Arc::new(super::tests::RejectingExec), + russh::client::Config::default(), + ) + .await; + let mut pty = client.channel_open_session().await.unwrap(); + pty.request_pty(true, "xterm", 80, 24, 0, 0, &[]) + .await + .unwrap(); + assert!(matches!(pty.wait().await, Some(ChannelMsg::Success))); + pty.request_subsystem(true, SUBSYSTEM).await.unwrap(); + assert!(matches!(pty.wait().await, Some(ChannelMsg::Failure))); + let mut upload = client.channel_open_session().await.unwrap(); + upload.request_subsystem(true, SUBSYSTEM).await.unwrap(); + assert!(matches!(upload.wait().await, Some(ChannelMsg::Success))); + let mut reader = Reader::new(); + assert!(matches!( + reader.frame(&mut upload).await, + Frame::Hello { .. } + )); + upload.exec(true, "cat").await.unwrap(); + assert!(matches!(upload.wait().await, Some(ChannelMsg::Failure))); + upload.request_shell(true).await.unwrap(); + assert!(matches!(upload.wait().await, Some(ChannelMsg::Failure))); + upload.request_subsystem(true, "sftp").await.unwrap(); + assert!(matches!(upload.wait().await, Some(ChannelMsg::Failure))); + upload + .request_pty(true, "xterm", 80, 24, 0, 0, &[]) + .await + .unwrap(); + assert!(matches!(upload.wait().await, Some(ChannelMsg::Failure))); + upload.close().await.unwrap(); + }) + .await + .expect("duplicate upload request disrupted SSH session"); +} + +struct DelayedExec { + session: Mutex>, + started: watch::Sender, + release: Mutex>>, +} + +#[async_trait::async_trait] +impl BoundaryExec for DelayedExec { + async fn exec(&self, _: ExecSpec) -> Result { + let release = self.release.lock().unwrap().take().unwrap(); + self.started.send_replace(true); + release.await.unwrap(); + Ok(self.session.lock().unwrap().take().unwrap()) + } +} + +#[tokio::test] +async fn cancellation_during_remote_startup_terminates_late_process() { + for disconnect in [false, true] { + tokio::time::timeout(DEADLINE, async { + let (session, control) = exec_session(); + let (release, receiver) = oneshot::channel(); + let started = watch::channel(false).0; + let client = test_client( + None, + Arc::new(DelayedExec { + session: Mutex::new(Some(session)), + started: started.clone(), + release: Mutex::new(Some(receiver)), + }), + russh::client::Config::default(), + ) + .await; + let mut channel = client.channel_open_session().await.unwrap(); + channel.request_subsystem(true, SUBSYSTEM).await.unwrap(); + assert!(matches!(channel.wait().await, Some(ChannelMsg::Success))); + let mut reader = Reader::new(); + assert!(matches!( + reader.frame(&mut channel).await, + Frame::Hello { .. } + )); + send( + &channel, + Frame::Begin(UploadPlan { + destination_directory: "/tmp/upload".into(), + expected_destination: None, + }), + ) + .await; + started + .subscribe() + .wait_for(|started| *started) + .await + .unwrap(); + if disconnect { + client + .disconnect(russh::Disconnect::ByApplication, "test", "") + .await + .unwrap(); + } else { + channel.close().await.unwrap(); + } + // A signal on another channel is a session-loop barrier proving + // that the close was handled while startup is still pending. + // Disconnect has no surviving channel, so wait for client closure. + if disconnect { + let _ = client.await; + } else { + let mut sibling = client.channel_open_session().await.unwrap(); + sibling + .request_subsystem(true, "unsupported") + .await + .unwrap(); + assert!(matches!(sibling.wait().await, Some(ChannelMsg::Failure))); + } + release.send(()).unwrap(); + control + .process + .terminated + .subscribe() + .wait_for(|terminated| *terminated) + .await + .unwrap(); + }) + .await + .expect("cancelled startup stranded a late extractor"); + } +} + +#[tokio::test] +async fn mismatched_end_and_extended_data_fail_without_result() { + for extended in [false, true] { + tokio::time::timeout(DEADLINE, async { + let (session, control) = exec_session(); + let client = test_client(None, Arc::new(Exec(Mutex::new([session].into()))), russh::client::Config::default()).await; + let (mut channel, mut reader) = start(&client).await; + if extended { + channel.extended_data(1, b"unexpected".as_slice()).await.unwrap(); + } else { + send(&channel, Frame::End(UploadSummary {bytes: 1, sha256: [0;32]})).await; + } + assert!(matches!(reader.frame(&mut channel).await, Frame::Error(failure) if failure.code == FailureCode::Protocol)); + control.process.terminated.subscribe().wait_for(|terminated| *terminated).await.unwrap(); + }).await.expect("invalid upload input did not terminate"); + } +} + +struct LoginProbe(watch::Sender>); + +#[async_trait::async_trait] +impl BoundaryExec for LoginProbe { + async fn exec(&self, spec: ExecSpec) -> Result { + self.0.send_replace(Some(spec.shell.unwrap().login)); + Err(BackendError::Process("login selection observed".into())) + } +} + +#[tokio::test] +async fn upload_preserves_login_shell_selection() { + for no_login in [false, true] { + tokio::time::timeout(DEADLINE, async { + let login = watch::channel(None).0; + let client = test_client( + None, + Arc::new(LoginProbe(login.clone())), + russh::client::Config::default(), + ) + .await; + let mut channel = client.channel_open_session().await.unwrap(); + if no_login { + channel + .set_env(true, "OPENSHELL_NO_LOGIN_SHELL", "1") + .await + .unwrap(); + assert!(matches!(channel.wait().await, Some(ChannelMsg::Success))); + } + channel.request_subsystem(true, SUBSYSTEM).await.unwrap(); + assert!(matches!(channel.wait().await, Some(ChannelMsg::Success))); + let mut reader = Reader::new(); + assert!(matches!( + reader.frame(&mut channel).await, + Frame::Hello { .. } + )); + send( + &channel, + Frame::Begin(UploadPlan { + destination_directory: "/tmp/upload".into(), + expected_destination: None, + }), + ) + .await; + assert!(matches!(reader.frame(&mut channel).await, Frame::Error(_))); + assert_eq!(*login.borrow(), Some(!no_login)); + }) + .await + .expect("login shell selection not applied to upload"); + } +} diff --git a/docs/how-it-works/sandboxes/overview.mdx b/docs/how-it-works/sandboxes/overview.mdx index 2e824c730e..2094a7280e 100644 --- a/docs/how-it-works/sandboxes/overview.mdx +++ b/docs/how-it-works/sandboxes/overview.mdx @@ -852,6 +852,16 @@ Upload files from your host into the sandbox: openshell sandbox upload my-sandbox ./src ``` +Uploads stream with bounded buffering and wait for the sandbox to consume +data, so the 4 MiB process-input budget does not limit the total upload size. +Upgrade both the CLI and the supervisor runtime used by your sandboxes to use +this behavior. The CLI reports an unsupported upload protocol before sending +file data to an older supervisor. + +Uploads extract directly into the destination. If you cancel an upload or it +fails, files already extracted remain and may be incomplete. Retry against the +existing sandbox after checking the destination, or remove the partial files. + When you omit the destination, OpenShell discovers the sandbox's working directory and uploads there. For a named local directory, OpenShell preserves the basename, matching `scp -r` and `cp -r`. If that directory already exists, diff --git a/e2e/rust/e2e-podman.sh b/e2e/rust/e2e-podman.sh index 0ebdc2359e..cd8b52e695 100755 --- a/e2e/rust/e2e-podman.sh +++ b/e2e/rust/e2e-podman.sh @@ -41,6 +41,7 @@ PODMAN_CI_TESTS=( podman_host_gateway podman_oci_identity provider_token_exchange + upload_backpressure ) if [ -n "${E2E_TEST}" ] && [ -n "${E2E_TEST_SET}" ]; then diff --git a/e2e/rust/tests/upload_backpressure.rs b/e2e/rust/tests/upload_backpressure.rs new file mode 100644 index 0000000000..e766303c37 --- /dev/null +++ b/e2e/rust/tests/upload_backpressure.rs @@ -0,0 +1,121 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![cfg(all(feature = "e2e", unix))] + +//! Regression for large built-in uploads through the CLI, gateway, supervisor, +//! and sandbox. A fast sender previously overflowed the supervisor's 4 MiB +//! pending-stdin budget while tar was creating many files (issue #4430). + +use std::fmt::Write as _; +use std::fs; +use std::os::unix::fs::{PermissionsExt, symlink}; +use std::time::Duration; + +use openshell_e2e::harness::binary::openshell_cmd; +use openshell_e2e::harness::output::strip_ansi; +use openshell_e2e::harness::sandbox::SandboxGuard; +use rand::{RngCore, SeedableRng, rngs::StdRng}; +use sha2::{Digest, Sha256}; +use tokio::time::timeout; + +const FILE_COUNT: usize = 12_000; +const FILE_BYTES: usize = 8_192; +const TRANSFER_TIMEOUT: Duration = Duration::from_secs(300); + +/// An archive substantially larger than the stdin budget must finish without +/// pacing at the caller, preserving every file's content and tar semantics. +#[tokio::test] +async fn large_upload_preserves_content_and_metadata_outside_workdir() { + let tmpdir = tempfile::tempdir().expect("create upload fixture tempdir"); + let source = tmpdir.path().join("large-project"); + fs::create_dir_all(source.join("payload")).expect("create payload directory"); + + // Distinct deterministic pseudorandom bytes keep the fixture reproducible + // without letting compression or identical-file checks hide dropped data. + let mut rng = StdRng::seed_from_u64(4430); + let mut payload = [0_u8; FILE_BYTES]; + let mut checksums = String::new(); + for index in 0..FILE_COUNT { + rng.fill_bytes(&mut payload); + let relative = format!("payload/{index:05}.bin"); + fs::write(source.join(&relative), payload).expect("write payload file"); + writeln!(checksums, "{:x} {relative}", Sha256::digest(payload)) + .expect("append file checksum"); + } + let manifest_digest = format!("{:x}", Sha256::digest(checksums.as_bytes())); + fs::write(source.join("SHA256SUMS"), checksums).expect("write checksum manifest"); + fs::write( + source.join("run.sh"), + "#!/bin/sh\nprintf '%s\\n' upload-executable-ok\n", + ) + .expect("write executable fixture"); + fs::set_permissions(source.join("run.sh"), fs::Permissions::from_mode(0o751)) + .expect("set executable mode"); + symlink("payload/00000.bin", source.join("payload-link")).expect("create relative symlink"); + symlink("missing-target", source.join("dangling-link")).expect("create dangling symlink"); + fs::write(source.join("conflict.txt"), "local-conflict").expect("write merge fixture"); + + let mut sandbox = SandboxGuard::create(&[]) + .await + .expect("create persistent upload sandbox"); + sandbox + .exec(&[ + "sh", + "-c", + "set -eu; test \"$(pwd -P)\" != /tmp/openshell-upload-backpressure; \ + mkdir -p /tmp/openshell-upload-backpressure/large-project; \ + printf remote-conflict > /tmp/openshell-upload-backpressure/large-project/conflict.txt; \ + printf remote-preserved > /tmp/openshell-upload-backpressure/large-project/unrelated.txt", + ]) + .await + .expect("prepare destination outside the canonical workdir"); + + // No sleep, bandwidth limit, or retry: the protocol must backpressure the + // normal CLI tar stream, even when the local sender outruns extraction. + let mut upload = openshell_cmd(); + upload + .env("RUST_LOG", "warn") + .args(["sandbox", "upload", &sandbox.name]) + .arg(&source) + .args(["/tmp/openshell-upload-backpressure", "--no-git-ignore"]); + let result = timeout(TRANSFER_TIMEOUT, upload.output()) + .await + .expect("large upload must not hang waiting for credit") + .expect("start large upload"); + let output = format!( + "{}{}", + String::from_utf8_lossy(&result.stdout), + String::from_utf8_lossy(&result.stderr) + ); + assert!(result.status.success(), "large upload failed: {output}"); + assert!( + !output.contains("sandbox sync operation failed; retrying"), + "large upload required a retry: {output}" + ); + + let verify = format!( + "set -eu; cd /tmp/openshell-upload-backpressure/large-project; \ + test \"$(find payload -type f | wc -l)\" -eq {FILE_COUNT}; \ + test \"$(sha256sum SHA256SUMS | cut -d ' ' -f 1)\" = {manifest_digest}; \ + sha256sum --check --quiet SHA256SUMS; \ + test \"$(stat -c %a run.sh)\" = 751; \ + test \"$(./run.sh)\" = upload-executable-ok; \ + test -L payload-link; test \"$(readlink payload-link)\" = payload/00000.bin; \ + test -L dangling-link; test ! -e dangling-link; \ + test \"$(readlink dangling-link)\" = missing-target; \ + test \"$(cat conflict.txt)\" = local-conflict; \ + test \"$(cat unrelated.txt)\" = remote-preserved; \ + printf large-upload-content-and-metadata-ok" + ); + let output = timeout(TRANSFER_TIMEOUT, sandbox.exec(&["sh", "-c", &verify])) + .await + .expect("large upload verification must not hang") + .expect("every uploaded file and metadata item must match"); + assert!( + strip_ansi(&output).contains("large-upload-content-and-metadata-ok"), + "expected complete upload verification marker:\n{output}" + ); + + sandbox.cleanup().await; +} diff --git a/skills/openshell-cli/SKILL.md b/skills/openshell-cli/SKILL.md index 7bfb049f49..df2145dc5f 100644 --- a/skills/openshell-cli/SKILL.md +++ b/skills/openshell-cli/SKILL.md @@ -410,6 +410,8 @@ openshell sandbox upload my-sandbox ./src openshell sandbox download my-sandbox output ./local-output ``` +Uploads use bounded streaming and wait for sandbox input consumption; the 4 MiB process-input budget is not a total upload-size limit. Upgrade both the CLI and the supervisor runtime used by the sandbox. An older supervisor rejects the upload protocol before file data is sent. Failed or canceled uploads may leave incomplete files in the destination; check it before retrying or remove the partial files. + Uploads inside a Git work tree honor `.gitignore` by default and stop if Git filtering fails or selects no files. Sources confirmed to be outside a Git work tree upload without filtering, with a warning that `.gitignore` rules are not applied. Git must be available to determine whether filtering applies; broken or inaccessible repositories stop the upload. Add `--no-git-ignore` for an intentional unfiltered upload. This also applies to `sandbox create --upload`. `sandbox create --upload` checks every upload locally before provisioning and creates no sandbox if one is rejected. If a transfer fails after provisioning, the sandbox remains running and earlier uploads may have completed. Retry with `sandbox upload` against that sandbox, or remove it with `sandbox delete`. From 89fb703fb86a23d0501f642da206ae36b3c6b863 Mon Sep 17 00:00:00 2001 From: Shiju Date: Tue, 13 Oct 2026 03:28:26 +0530 Subject: [PATCH 2/5] fix(cli): preserve upload cancellation and errors Keep cancellation active through upload setup and retry delays, and terminate cancelled SSH probes. Preserve protocol failures through output completion and report captured diagnostics on early extraction failure. Remove redundant upload validation and extend the affected regressions. Signed-off-by: Shiju --- crates/openshell-cli/src/ssh.rs | 173 ++++++++- crates/openshell-cli/src/ssh/upload.rs | 25 +- crates/openshell-core/src/upload.rs | 7 - .../src/ssh/upload.rs | 342 ++++++++++-------- .../src/ssh/upload_tests.rs | 66 +++- 5 files changed, 415 insertions(+), 198 deletions(-) diff --git a/crates/openshell-cli/src/ssh.rs b/crates/openshell-cli/src/ssh.rs index 5de7a0686f..1e50ef574e 100644 --- a/crates/openshell-cli/src/ssh.rs +++ b/crates/openshell-cli/src/ssh.rs @@ -1502,7 +1502,7 @@ pub async fn sandbox_sync_up_files( files: files.to_vec(), archive_prefix: file_list_archive_prefix(local_path), }; - retry_sandbox_sync("upload", || { + retry_sandbox_upload(|| { let source = source.clone(); async move { ssh_tar_upload(server, name, dest, source, tls, workspace).await } }) @@ -1542,7 +1542,7 @@ pub async fn sandbox_sync_up( local_path: local_path.to_path_buf(), tar_name, }; - retry_sandbox_sync("upload", || { + retry_sandbox_upload(|| { let source = source.clone(); async move { ssh_tar_upload(server, name, sandbox_path, source, tls, workspace).await } }) @@ -1589,9 +1589,10 @@ async fn ssh_run_capture_stdout(session: &SshSessionConfig, command: &str) -> Re .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::inherit()); - let output = tokio::task::spawn_blocking(move || ssh.output()) + let output = TokioCommand::from(ssh) + .kill_on_drop(true) + .output() .await - .into_diagnostic()? .into_diagnostic()?; if !output.status.success() { return Err(miette::miette!( @@ -1711,6 +1712,25 @@ async fn sandbox_sync_down_once( } } +async fn retry_sandbox_upload(run: F) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + #[cfg(unix)] + { + // Tokio retains handlers after listeners drop; cover setup and retry delays too. + let mut signals = TerminationSignals::new()?; + tokio::select! { + biased; + _ = signals.recv() => Err(Report::new(upload::UploadError::Cancelled)), + result = retry_sandbox_sync("upload", run) => result, + } + } + #[cfg(not(unix))] + retry_sandbox_sync("upload", run).await +} + async fn retry_sandbox_sync(operation: &str, mut run: F) -> Result where F: FnMut() -> Fut, @@ -2498,6 +2518,151 @@ mod tests { ))); } + #[cfg(unix)] + #[test] + fn upload_cancellation_covers_retry_delay_and_setup() { + use std::os::unix::fs::PermissionsExt; + + const PHASE_ENV: &str = "OPENSHELL_UPLOAD_CANCEL_PHASE"; + const MARKER_ENV: &str = "OPENSHELL_UPLOAD_CANCEL_MARKER"; + const PID_ENV: &str = "OPENSHELL_UPLOAD_CANCEL_PROBE_PID"; + struct Probe(std::process::Child); + impl Drop for Probe { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } + } + if let Some(phase) = std::env::var_os(PHASE_ENV) { + let marker = PathBuf::from(std::env::var_os(MARKER_ENV).unwrap()); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(async { + let mut attempts = 0; + let error = retry_sandbox_upload(|| { + attempts += 1; + assert_eq!(attempts, 1, "cancelled upload started another attempt"); + let marker = marker.clone(); + let setup = phase == "setup"; + async move { + if setup { + let session = SshSessionConfig { + proxy_command: String::new(), + sandbox_id: String::new(), + sandbox_name: String::new(), + gateway_url: String::new(), + token: String::new(), + main_terminal: false, + replica: None, + }; + return ssh_run_capture_stdout(&session, "unused").await.map(|_| ()); + } + let mut command = Command::new("sh"); + command.args(["-c", "exit 255"]); + let error = upload::run( + command, + UploadPlan { + destination_directory: "/tmp/upload".into(), + expected_destination: None, + }, + |_| Ok(()), + ) + .await + .unwrap_err(); + assert!(matches!(error, upload::UploadError::Transport(_))); + // A spawned task runs only after the current-thread + // runtime reaches the retry wait. + tokio::spawn(async move { + fs::write(marker, "retry delay").unwrap(); + }); + Err(Report::new(error)) + } + }) + .await + .unwrap_err(); + assert!(matches!( + error.downcast_ref::(), + Some(upload::UploadError::Cancelled) + )); + assert_eq!(attempts, 1); + if phase == "setup" { + let pid = fs::read_to_string(std::env::var_os(PID_ENV).unwrap()) + .unwrap() + .parse::() + .unwrap(); + tokio::time::timeout(Duration::from_millis(750), async { + while nix::sys::signal::kill(nix::unistd::Pid::from_raw(pid), None).is_ok() + { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("cancelled probe must be terminated and reaped"); + } + }); + return; + } + + for phase in ["retry delay", "setup"] { + let temporary = tempfile::tempdir().unwrap(); + let marker = temporary.path().join("ready"); + let mut command = Command::new(std::env::current_exe().unwrap()); + command + .args([ + "--exact", + "ssh::tests::upload_cancellation_covers_retry_delay_and_setup", + "--nocapture", + ]) + .env(PHASE_ENV, phase) + .env(MARKER_ENV, &marker) + .env(PID_ENV, temporary.path().join("probe-pid")) + .stdin(Stdio::null()) + .stdout(Stdio::null()); + if phase == "setup" { + let ssh = temporary.path().join("ssh"); + fs::write( + &ssh, + "#!/bin/sh\n\ + printf '%s' \"$$\" > \"$OPENSHELL_UPLOAD_CANCEL_PROBE_PID\"\n\ + printf setup > \"$OPENSHELL_UPLOAD_CANCEL_MARKER\"\n\ + exec sleep 60\n", + ) + .unwrap(); + fs::set_permissions(ssh, fs::Permissions::from_mode(0o755)).unwrap(); + let mut path = temporary.path().as_os_str().to_os_string(); + path.push(":"); + path.push(std::env::var_os("PATH").unwrap_or_default()); + command.env("PATH", path); + } + let mut probe = Probe(command.spawn().unwrap()); + let deadline = Instant::now() + Duration::from_secs(5); + while !fs::read_to_string(&marker).is_ok_and(|ready| ready == phase) { + assert!(probe.0.try_wait().unwrap().is_none(), "probe exited early"); + assert!(Instant::now() < deadline, "probe did not reach {phase}"); + std::thread::sleep(Duration::from_millis(10)); + } + nix::sys::signal::kill( + nix::unistd::Pid::from_raw(probe.0.id().cast_signed()), + Signal::SIGINT, + ) + .unwrap(); + let deadline = Instant::now() + Duration::from_secs(1); + loop { + if let Some(status) = probe.0.try_wait().unwrap() { + assert!(status.success(), "upload cancellation failed: {status}"); + break; + } + assert!( + Instant::now() < deadline, + "cancellation waited through {phase}" + ); + std::thread::sleep(Duration::from_millis(10)); + } + } + } + #[test] fn connect_recovery_starts_only_for_established_transport_failures() { assert!(should_recover_connect( diff --git a/crates/openshell-cli/src/ssh/upload.rs b/crates/openshell-cli/src/ssh/upload.rs index 247e4e09c9..e4400eebfd 100644 --- a/crates/openshell-cli/src/ssh/upload.rs +++ b/crates/openshell-cli/src/ssh/upload.rs @@ -9,8 +9,7 @@ use std::time::Duration; use miette::Diagnostic; use openshell_core::upload::{ - Frame, MAX_DATA_BYTES, MAX_DIAGNOSTIC_BYTES, ProtocolError, UploadPlan, UploadSummary, - WINDOW_BYTES, read_frame, write_frame, + Frame, MAX_DIAGNOSTIC_BYTES, ProtocolError, UploadPlan, UploadSummary, read_frame, write_frame, }; use sha2::{Digest, Sha256}; use thiserror::Error; @@ -120,15 +119,6 @@ where let diagnostics = tokio::spawn(read_diagnostic(stderr)); let diagnostic_guard = AbortTasks(vec![diagnostics.abort_handle()]); - #[cfg(unix)] - let mut signals = super::TerminationSignals::new() - .map_err(|error| UploadError::Transport(error.to_string()))?; - #[cfg(unix)] - let result = tokio::select! { - result = transfer(stdout, stdin, plan, producer) => result, - _ = signals.recv() => Err(UploadError::Cancelled), - }; - #[cfg(not(unix))] let result = transfer(stdout, stdin, plan, producer).await; let status = if result.is_err() { @@ -211,17 +201,6 @@ async fn negotiate( else { return Err(server_failure(hello, "expected HELLO")); }; - // A peer cannot increase local memory use by advertising larger limits. - if window_bytes == 0 - || window_bytes > WINDOW_BYTES - || max_data_bytes == 0 - || usize::try_from(max_data_bytes).unwrap_or(usize::MAX) > MAX_DATA_BYTES - || max_data_bytes > window_bytes - { - return Err(UploadError::Protocol( - "supervisor advertised invalid upload limits".to_string(), - )); - } write_frame(writer, &Frame::Begin(plan)) .await .map_err(protocol_error)?; @@ -433,7 +412,7 @@ where #[cfg(test)] mod tests { use super::*; - use openshell_core::upload::{FailureCode, UploadFailure}; + use openshell_core::upload::{FailureCode, MAX_DATA_BYTES, UploadFailure, WINDOW_BYTES}; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use tokio::io::{DuplexStream, ReadHalf, WriteHalf}; diff --git a/crates/openshell-core/src/upload.rs b/crates/openshell-core/src/upload.rs index c7969f9b06..588159e0f0 100644 --- a/crates/openshell-core/src/upload.rs +++ b/crates/openshell-core/src/upload.rs @@ -33,13 +33,6 @@ pub enum ProtocolError { Json(#[from] serde_json::Error), } -impl ProtocolError { - #[must_use] - pub const fn is_io(&self) -> bool { - matches!(self, Self::Io(_)) - } -} - #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum DestinationKind { diff --git a/crates/openshell-supervisor-process/src/ssh/upload.rs b/crates/openshell-supervisor-process/src/ssh/upload.rs index cb0453051d..21fbc532c7 100644 --- a/crates/openshell-supervisor-process/src/ssh/upload.rs +++ b/crates/openshell-supervisor-process/src/ssh/upload.rs @@ -269,13 +269,23 @@ async fn finish( handle: &Handle, result: Result, ) { + let (result, process) = shared.state.lock().map_or_else( + |_| { + ( + Err(failure(FailureCode::Internal, "upload state unavailable")), + None, + ) + }, + |mut state| { + // Commit the result under the same lock that records protocol failure. + let result = state.failure.take().map_or(result, Err); + state.phase = Phase::Finished; + state.sender.take(); + state.input_task.take(); + (result, state.process.take()) + }, + ); let failed = result.is_err(); - let process = shared.state.lock().ok().and_then(|mut state| { - state.phase = Phase::Finished; - state.sender.take(); - state.input_task.take(); - state.process.take() - }); if failed && let Some(process) = process { tokio::spawn(async move { let _ = process.terminate().await; @@ -342,186 +352,202 @@ async fn drive( } changed.await; }; - let command = build_extract_command(&plan) - .map_err(|error| failure(FailureCode::Protocol, error.to_string()))?; - // Remote backends may admit a process before returning its attachment. - // Keep startup owned independently of this channel: cancellation must - // terminate even a result delivered after the upload owner is dropped. - let (started_tx, started_rx) = oneshot::channel(); - tokio::spawn(async move { - let result = boundary_exec - .exec(ExecSpec { - program: String::new(), - args: Vec::new(), - shell: Some(ShellSpec { - command: Some(command), - login, - }), - runtime_helper: None, - env: Vec::new(), - workdir: None, - pty: false, - }) + let extraction = async { + let command = build_extract_command(&plan) + .map_err(|error| failure(FailureCode::Protocol, error.to_string()))?; + // Remote backends may admit a process before returning its attachment. + // Keep startup owned independently of this channel: cancellation must + // terminate even a result delivered after the upload owner is dropped. + let (started_tx, started_rx) = oneshot::channel(); + tokio::spawn(async move { + let result = boundary_exec + .exec(ExecSpec { + program: String::new(), + args: Vec::new(), + shell: Some(ShellSpec { + command: Some(command), + login, + }), + runtime_helper: None, + env: Vec::new(), + workdir: None, + pty: false, + }) + .await + .map(|exec| PendingExec(Some(exec))); + // A queued PendingExec also terminates on drop if the receiver was + // cancelled after send succeeded but before receiving the attachment. + let _ = started_tx.send(result); + }); + let mut pending = started_rx .await - .map(|exec| PendingExec(Some(exec))); - // A queued PendingExec also terminates on drop if the receiver was - // cancelled after send succeeded but before receiving the attachment. - let _ = started_tx.send(result); - }); - let mut pending = started_rx - .await - .map_err(|_| failure(FailureCode::Internal, "extraction startup stopped"))? - .map_err(|error| failure(FailureCode::Extract, error.to_string()))?; - { - let mut state = shared - .state - .lock() - .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; - // Channel drop can run on another worker after the receive completes. - // Keep the startup guard armed until ownership is registered under the - // same lock that cancellation uses to mark the channel finished. - if state.phase == Phase::Finished || state.failure.is_some() { - return Err(failure( - FailureCode::Protocol, - "upload cancelled during startup", - )); - } - state.process = Some( - pending - .0 - .as_ref() - .expect("owned startup result") - .process - .clone(), - ); - } - let mut exec = pending.0.take().expect("registered startup result"); - let stdin = exec - .stdin - .take() - .ok_or_else(|| failure(FailureCode::Internal, "extraction process has no stdin"))?; - let (sender, input_task, mut progress) = InputTask::spawn_with_progress(stdin); - { - let mut state = shared - .state - .lock() - .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; - if state.phase == Phase::Finished { - return Err(failure( - FailureCode::Protocol, - "upload cancelled during startup", - )); - } - if let Some(error) = state.failure.take() { - return Err(error); + .map_err(|_| failure(FailureCode::Internal, "extraction startup stopped"))? + .map_err(|error| failure(FailureCode::Extract, error.to_string()))?; + { + let mut state = shared + .state + .lock() + .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; + // Channel drop can run on another worker after the receive completes. + // Keep the startup guard armed until ownership is registered under the + // same lock that cancellation uses to mark the channel finished. + if state.phase == Phase::Finished || state.failure.is_some() { + return Err(failure( + FailureCode::Protocol, + "upload cancelled during startup", + )); + } + state.process = Some( + pending + .0 + .as_ref() + .expect("owned startup result") + .process + .clone(), + ); } - state.sender = Some(sender); - state.input_task = Some(input_task); - state.progress = Some(progress.clone()); - state.phase = Phase::Receiving; - } - let diagnostic = Arc::new(Mutex::new(Vec::new())); - let stdout = DrainTask::spawn(exec.stdout, diagnostic.clone()); - let stderr = exec - .stderr - .map(|stderr| DrainTask::spawn(stderr, diagnostic.clone())); - let mut exit = Box::pin(exec.process.wait()); - if !send(handle, channel, Frame::Ready).await { - return Err(failure( - FailureCode::Internal, - "upload reply delivery failed", - )); - } - let mut acknowledged = 0; - let status = loop { - let changed = shared.changed.notified(); + let mut exec = pending.0.take().expect("registered startup result"); + let stdin = exec + .stdin + .take() + .ok_or_else(|| failure(FailureCode::Internal, "extraction process has no stdin"))?; + let (sender, input_task, mut progress) = InputTask::spawn_with_progress(stdin); { let mut state = shared .state .lock() .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; + if state.phase == Phase::Finished { + return Err(failure( + FailureCode::Protocol, + "upload cancelled during startup", + )); + } if let Some(error) = state.failure.take() { return Err(error); } + state.sender = Some(sender); + state.input_task = Some(input_task); + state.progress = Some(progress.clone()); + state.phase = Phase::Receiving; } - let latest = *progress.borrow_and_update(); - if latest.state == InputState::Failed { + let diagnostic = Arc::new(Mutex::new(Vec::new())); + let stdout = DrainTask::spawn(exec.stdout, diagnostic.clone()); + let stderr = exec + .stderr + .map(|stderr| DrainTask::spawn(stderr, diagnostic.clone())); + let mut exit = Box::pin(exec.process.wait()); + if !send(handle, channel, Frame::Ready).await { return Err(failure( - FailureCode::InputWrite, - "failed to write extraction stdin", + FailureCode::Internal, + "upload reply delivery failed", )); } - if latest.consumed - acknowledged >= ACK_INTERVAL - || (latest.state == InputState::Drained && latest.consumed != acknowledged) - { - if !send(handle, channel, Frame::Ack(latest.consumed)).await { + let mut acknowledged = 0; + let status = loop { + let latest = *progress.borrow_and_update(); + if latest.state == InputState::Failed { return Err(failure( - FailureCode::Internal, - "upload reply delivery failed", + FailureCode::InputWrite, + diagnostic_message(&diagnostic, "failed to write extraction stdin"), )); } - acknowledged = latest.consumed; - } - tokio::select! { - status = &mut exit => { - let status = status.map_err(|error| failure(FailureCode::Extract, error.to_string()))?; - let final_progress = *progress.borrow(); - let (final_summary, final_received) = shared.state.lock().map(|state| (state.summary, state.received)).map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; - if final_summary.is_none() || final_progress.consumed != final_received { - return Err(extract_failure(status, "extraction process exited before upload completed")); - } - // A pipe shutdown can wake the child just before the writer - // publishes Drained. Await that final watch update rather than - // treating the successful EOF race as an early extraction exit. - while progress.borrow().state == InputState::Writing { - progress.changed().await.map_err(|_| failure(FailureCode::InputWrite, "extraction stdin writer stopped"))?; - } - if progress.borrow().state != InputState::Drained { - return Err(failure(FailureCode::InputWrite, "failed to close extraction stdin")); - } - if final_received != acknowledged && !send(handle, channel, Frame::Ack(final_received)).await { - return Err(failure(FailureCode::Internal, "upload reply delivery failed")); + if latest.consumed - acknowledged >= ACK_INTERVAL + || (latest.state == InputState::Drained && latest.consumed != acknowledged) + { + if !send(handle, channel, Frame::Ack(latest.consumed)).await { + return Err(failure( + FailureCode::Internal, + "upload reply delivery failed", + )); } - break status; + acknowledged = latest.consumed; } - () = changed => {} - result = progress.changed() => { - if result.is_err() { - return Err(failure(FailureCode::InputWrite, "extraction stdin writer stopped")); + tokio::select! { + status = &mut exit => { + let status = status.map_err(|error| failure(FailureCode::Extract, error.to_string()))?; + let final_progress = *progress.borrow(); + let (final_summary, final_received) = shared.state.lock().map(|state| (state.summary, state.received)).map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; + if final_summary.is_none() || final_progress.consumed != final_received { + return Err(extract_failure(status, &diagnostic_message(&diagnostic, "extraction process exited before upload completed"))); + } + // A pipe shutdown can wake the child just before the writer + // publishes Drained. Await that final watch update rather than + // treating the successful EOF race as an early extraction exit. + while progress.borrow().state == InputState::Writing { + progress.changed().await.map_err(|_| failure(FailureCode::InputWrite, "extraction stdin writer stopped"))?; + } + if progress.borrow().state != InputState::Drained { + return Err(failure(FailureCode::InputWrite, diagnostic_message(&diagnostic, "failed to close extraction stdin"))); + } + if final_received != acknowledged && !send(handle, channel, Frame::Ack(final_received)).await { + return Err(failure(FailureCode::Internal, "upload reply delivery failed")); + } + break status; + } + result = progress.changed() => { + if result.is_err() { + return Err(failure(FailureCode::InputWrite, "extraction stdin writer stopped")); + } } } - } - }; - stdout - .join() - .await - .map_err(|error| failure(FailureCode::Extract, error.to_string()))?; - if let Some(stderr) = stderr { - stderr + }; + stdout .join() .await .map_err(|error| failure(FailureCode::Extract, error.to_string()))?; - } - let status = if let Some(output_status) = exec.output_status { - output_status - .await - .map_err(|_| failure(FailureCode::Extract, "extraction output delivery failed"))? - } else { - status + if let Some(stderr) = stderr { + stderr + .join() + .await + .map_err(|error| failure(FailureCode::Extract, error.to_string()))?; + } + let status = if let Some(output_status) = exec.output_status { + output_status + .await + .map_err(|_| failure(FailureCode::Extract, "extraction output delivery failed"))? + } else { + status + }; + if status != BoundaryExitStatus::Exited(0) { + let diagnostic = diagnostic_message(&diagnostic, "extraction process failed"); + return Err(extract_failure(status, &diagnostic)); + } + shared + .state + .lock() + .ok() + .and_then(|state| state.summary) + .ok_or_else(|| failure(FailureCode::Internal, "upload summary unavailable")) }; - if status != BoundaryExitStatus::Exited(0) { - let diagnostic = diagnostic.lock().map_or_else( - |_| String::new(), - |bytes| String::from_utf8_lossy(&bytes).into_owned(), - ); - return Err(extract_failure(status, &diagnostic)); + tokio::pin!(extraction); + loop { + let changed = shared.changed.notified(); + let error = shared + .state + .lock() + .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))? + .failure + .clone(); + if let Some(error) = error { + return Err(error); + } + tokio::select! { + result = &mut extraction => return result, + () = changed => {} + } } - shared - .state +} + +fn diagnostic_message(diagnostic: &Mutex>, fallback: &str) -> String { + diagnostic .lock() .ok() - .and_then(|state| state.summary) - .ok_or_else(|| failure(FailureCode::Internal, "upload summary unavailable")) + .filter(|bytes| !bytes.is_empty()) + .map_or_else( + || fallback.to_string(), + |bytes| String::from_utf8_lossy(&bytes).into_owned(), + ) } fn extract_failure(status: BoundaryExitStatus, diagnostic: &str) -> UploadFailure { diff --git a/crates/openshell-supervisor-process/src/ssh/upload_tests.rs b/crates/openshell-supervisor-process/src/ssh/upload_tests.rs index 31e75b8e2c..f9cd4c4ef5 100644 --- a/crates/openshell-supervisor-process/src/ssh/upload_tests.rs +++ b/crates/openshell-supervisor-process/src/ssh/upload_tests.rs @@ -23,6 +23,7 @@ const DEADLINE: Duration = Duration::from_secs(10); struct Process { status: watch::Sender>, + waited: watch::Sender, terminated: watch::Sender, signals: mpsc::UnboundedSender, } @@ -31,13 +32,15 @@ struct Process { impl BoundaryProcess for Process { async fn wait(&self) -> Result { let mut status = self.status.subscribe(); - Ok(status + let result = status .wait_for(Option::is_some) .await .unwrap() .as_ref() .copied() - .unwrap()) + .unwrap(); + self.waited.send_replace(true); + Ok(result) } async fn signal(&self, signal: BoundarySignal) -> Result<(), BackendError> { @@ -82,6 +85,7 @@ fn exec_session() -> (ExecSession, Control) { let (signal_tx, signals) = mpsc::unbounded_channel(); let process = Arc::new(Process { status: watch::channel(None).0, + waited: watch::channel(false).0, terminated: watch::channel(false).0, signals: signal_tx, }); @@ -311,7 +315,7 @@ async fn credit_violation_and_truncated_eof_are_typed_and_terminate() { async fn broken_stdin_and_early_exit_wake_credit_waiter() { for early_exit in [false, true] { tokio::time::timeout(DEADLINE, async { - let (session, control) = exec_session(); + let (session, mut control) = exec_session(); let client = test_client( None, Arc::new(Exec(Mutex::new([session].into()))), @@ -319,6 +323,11 @@ async fn broken_stdin_and_early_exit_wake_credit_waiter() { ) .await; let (mut channel, mut reader) = start(&client).await; + let diagnostic = "mkdir: /tmp/existing-file/child: Not a directory"; + let mut stderr = diagnostic.as_bytes().to_vec(); + stderr.resize(8192, b' '); + // Filling more than the pipe holds proves the drain captured the cause. + control.stderr.write_all(&stderr).await.unwrap(); if early_exit { control.process.status.send_replace(Some(BoundaryExitStatus::Exited(1))); } else { @@ -326,7 +335,7 @@ async fn broken_stdin_and_early_exit_wake_credit_waiter() { send(&channel, Frame::Data(vec![7; MAX_DATA_BYTES])).await; } let expected = if early_exit { FailureCode::Extract } else { FailureCode::InputWrite }; - assert!(matches!(reader.frame(&mut channel).await, Frame::Error(error) if error.code == expected)); + assert!(matches!(reader.frame(&mut channel).await, Frame::Error(error) if error.code == expected && error.message.contains(diagnostic))); control.process.terminated.subscribe().wait_for(|terminated| *terminated).await.unwrap(); }) .await @@ -505,11 +514,38 @@ async fn cancellation_during_remote_startup_terminates_late_process() { #[tokio::test] async fn mismatched_end_and_extended_data_fail_without_result() { - for extended in [false, true] { + for (extended, after_exit, pending_output_status) in [ + (false, false, false), + (true, false, false), + (true, true, false), + (true, true, true), + ] { tokio::time::timeout(DEADLINE, async { - let (session, control) = exec_session(); + let (mut session, mut control) = exec_session(); + let (output_status, receiver) = oneshot::channel(); + if pending_output_status { + session.output_status = Some(receiver); + } let client = test_client(None, Arc::new(Exec(Mutex::new([session].into()))), russh::client::Config::default()).await; let (mut channel, mut reader) = start(&client).await; + if after_exit { + let summary = UploadSummary { bytes: 1, sha256: Sha256::digest(b"x").into() }; + send(&channel, Frame::Data(b"x".to_vec())).await; + send(&channel, Frame::End(summary)).await; + let mut bytes = Vec::new(); + control.stdin.read_to_end(&mut bytes).await.unwrap(); + assert_eq!(bytes, b"x"); + assert_eq!(reader.frame(&mut channel).await, Frame::Ack(1)); + control.process.status.send_replace(Some(BoundaryExitStatus::Exited(0))); + control.process.waited.subscribe().wait_for(|waited| *waited).await.unwrap(); + if pending_output_status { + control.stdout.shutdown().await.unwrap(); + control.stderr.shutdown().await.unwrap(); + assert_eq!(control.stdout.read(&mut [0]).await.unwrap(), 0); + assert_eq!(control.stderr.read(&mut [0]).await.unwrap(), 0); + tokio::task::yield_now().await; + } + } if extended { channel.extended_data(1, b"unexpected".as_slice()).await.unwrap(); } else { @@ -517,6 +553,24 @@ async fn mismatched_end_and_extended_data_fail_without_result() { } assert!(matches!(reader.frame(&mut channel).await, Frame::Error(failure) if failure.code == FailureCode::Protocol)); control.process.terminated.subscribe().wait_for(|terminated| *terminated).await.unwrap(); + // Error and cleanup must arrive while output EOF/status is still held. + if after_exit && !pending_output_status { + assert_eq!(control.stdout.read(&mut [0]).await.unwrap(), 0); + assert_eq!(control.stderr.read(&mut [0]).await.unwrap(), 0); + } + if pending_output_status { + assert!(output_status.is_closed()); + } + assert!(reader.frames.is_empty()); + while let Some(message) = channel.wait().await { + match message { + ChannelMsg::ExitStatus { exit_status } => reader.statuses.push(exit_status), + ChannelMsg::Data { .. } => panic!("unexpected upload frame after ERROR"), + ChannelMsg::Close => break, + _ => {} + } + } + assert_eq!(reader.statuses, [74]); }).await.expect("invalid upload input did not terminate"); } } From 3aeda041ac67d077364a909c4f28bfa8f3084246 Mon Sep 17 00:00:00 2001 From: Shiju Date: Tue, 13 Oct 2026 04:11:00 +0530 Subject: [PATCH 3/5] refactor(cli): carry upload credit over SSH streams Replace the framed upload protocol with SSH's own streams. The client sends one JSON plan line, the raw archive and EOF. The supervisor's stdout carries only cumulative totals of bytes written to extraction stdin, extraction output goes to stderr, and the exit status reports the result. Extraction starts inline through the ordinary exec path. This removes the frame codec, SHA-256 digests, the CLI sha2 dependency and the supervisor upload state machine while keeping the 1 MiB credit window, the 4 MiB input guard, cancellation and error reporting. A failed archive now stops SSH before closing stdin, so EOF never marks a truncated archive complete. Signed-off-by: Shiju --- Cargo.lock | 1 - crates/openshell-cli/Cargo.toml | 1 - crates/openshell-cli/src/ssh/upload.rs | 720 ++++++++---------- crates/openshell-core/src/upload.rs | 431 +++-------- crates/openshell-supervisor-process/README.md | 42 +- .../openshell-supervisor-process/src/ssh.rs | 187 ++++- .../src/ssh/upload.rs | 697 ++++------------- .../src/ssh/upload_tests.rs | 445 +++++------ 8 files changed, 922 insertions(+), 1602 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b44a5c2ad9..6a2c214e2c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4280,7 +4280,6 @@ dependencies = [ "rustls-pemfile", "serde", "serde_json", - "sha2 0.10.9", "tar", "temp-env", "tempfile", diff --git a/crates/openshell-cli/Cargo.toml b/crates/openshell-cli/Cargo.toml index a33a665201..5c92eefe5b 100644 --- a/crates/openshell-cli/Cargo.toml +++ b/crates/openshell-cli/Cargo.toml @@ -65,7 +65,6 @@ anyhow = { workspace = true } # File archiving (tar-over-SSH sync) tar = "0.4" tempfile = "3" -sha2 = { workspace = true } # OIDC/Auth oauth2 = { version = "5", default-features = false, features = ["reqwest"] } diff --git a/crates/openshell-cli/src/ssh/upload.rs b/crates/openshell-cli/src/ssh/upload.rs index e4400eebfd..8f3453ee7c 100644 --- a/crates/openshell-cli/src/ssh/upload.rs +++ b/crates/openshell-cli/src/ssh/upload.rs @@ -8,20 +8,19 @@ use std::process::{Command, ExitStatus, Stdio}; use std::time::Duration; use miette::Diagnostic; -use openshell_core::upload::{ - Frame, MAX_DIAGNOSTIC_BYTES, ProtocolError, UploadPlan, UploadSummary, read_frame, write_frame, -}; -use sha2::{Digest, Sha256}; +use openshell_core::upload::{DESTINATION_CHANGED_EXIT_CODE, UploadPlan, WINDOW_BYTES}; use thiserror::Error; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::process::{Child, Command as TokioCommand}; -use tokio::sync::mpsc; +use tokio::sync::{mpsc, watch}; use tokio::task::AbortHandle; const NEGOTIATION_TIMEOUT: Duration = Duration::from_secs(10); const CHILD_EXIT_TIMEOUT: Duration = Duration::from_secs(5); +const CHUNK_BYTES: usize = 64 * 1024; +const MAX_DIAGNOSTIC_BYTES: usize = 16 * 1024; /// The producer can hold at most two chunks in addition to its current buffer -/// and the chunk held by the driver. It cannot build an entire archive in RAM. +/// and the chunk held by the sender. It cannot build an entire archive in RAM. const PRODUCER_QUEUE_CHUNKS: usize = 2; #[derive(Debug, Error, Diagnostic)] @@ -92,7 +91,23 @@ impl Drop for AbortTasks { } } -/// Start the producer only after the subsystem has accepted the destination. +/// How a transfer ended before every archive byte was acknowledged. +#[derive(Debug)] +enum Stopped { + /// The SSH stream closed or failed; the exit status explains why. + Stream(String), + /// A local or protocol failure; the SSH child must be terminated. + Local(UploadError), +} + +/// Archive bytes sent and the final total the supervisor acknowledged. +#[derive(Debug)] +struct Sent { + bytes: u64, + acknowledged: u64, +} + +/// Start the producer only after the subsystem has started extraction. /// The archive callback uses the existing tar builder unchanged. pub(super) async fn run

( mut command: Command, @@ -102,7 +117,8 @@ pub(super) async fn run

( where P: FnOnce(&mut dyn Write) -> miette::Result<()> + Send + 'static, { - plan.validate() + let plan = plan + .encode() .map_err(|error| UploadError::Protocol(error.to_string()))?; command .stdin(Stdio::piped()) @@ -113,54 +129,66 @@ where .spawn() .map_err(|error| UploadError::Transport(error.to_string()))?; let mut child = ChildGuard(Some(child)); - let stdin = child.child().stdin.take().expect("piped SSH stdin"); + // EOF marks the end of the archive, so a failed transfer must stop SSH + // before its stdin closes. The transfer closes stdin only on success. + let mut stdin = Some(child.child().stdin.take().expect("piped SSH stdin")); let stdout = child.child().stdout.take().expect("piped SSH stdout"); let stderr = child.child().stderr.take().expect("piped SSH stderr"); let diagnostics = tokio::spawn(read_diagnostic(stderr)); let diagnostic_guard = AbortTasks(vec![diagnostics.abort_handle()]); - let result = transfer(stdout, stdin, plan, producer).await; - - let status = if result.is_err() { - child.finish(true).await - } else if let Ok(status) = tokio::time::timeout(CHILD_EXIT_TIMEOUT, child.finish(false)).await { - status - } else { + let transfer = match transfer(stdout, &mut stdin, plan, producer).await { + Ok(sent) => Ok(sent), + Err(Stopped::Stream(message)) => Err(message), + Err(Stopped::Local(error)) => { + let _ = child.finish(true).await; + return Err(error); + } + }; + drop(stdin); + let Ok(status) = tokio::time::timeout(CHILD_EXIT_TIMEOUT, child.finish(false)).await else { let _ = child.finish(true).await; - return Err(UploadError::Protocol( - "SSH did not exit after the upload result".to_string(), + return Err(UploadError::Transport( + "SSH did not exit after the upload stream ended".to_string(), )); }; + let status = status.map_err(|error| UploadError::Transport(error.to_string()))?; let diagnostic = tokio::time::timeout(CHILD_EXIT_TIMEOUT, diagnostics) .await .ok() .and_then(Result::ok) .unwrap_or_default(); drop(diagnostic_guard); - // OpenSSH reports subsystem rejection as status 255, like other SSH - // failures. Its explicit request diagnostic distinguishes an old server - // from a transient connection failure, without sending any archive bytes. - if matches!(result, Err(UploadError::Transport(_))) - && diagnostic.contains("subsystem request failed") - { - return Err(UploadError::Unsupported); - } - match result { - Err(UploadError::Transport(message)) if !diagnostic.trim().is_empty() => { - return Err(UploadError::Transport(format!("{message}: {diagnostic}"))); + let diagnostic = diagnostic.trim(); + match status.code() { + Some(0) => match transfer { + Ok(sent) if sent.acknowledged == sent.bytes => Ok(()), + Ok(sent) => Err(UploadError::Protocol(format!( + "sandbox acknowledged {} of {} upload bytes", + sent.acknowledged, sent.bytes + ))), + Err(message) => Err(UploadError::Protocol(format!( + "upload stream ended before completion: {message}" + ))), + }, + // OpenSSH reports subsystem rejection as status 255, like other SSH + // failures. Its explicit request diagnostic distinguishes an old server + // from a transient connection failure, without sending any archive bytes. + Some(super::SSH_TRANSPORT_FAILURE_EXIT_CODE) + if diagnostic.contains("subsystem request failed") => + { + Err(UploadError::Unsupported) } - result => result?, + Some(super::SSH_TRANSPORT_FAILURE_EXIT_CODE) | None => Err(UploadError::Transport( + format!("SSH exited with {status}: {diagnostic}"), + )), + Some(DESTINATION_CHANGED_EXIT_CODE) => Err(UploadError::Remote(format!( + "upload destination changed type before extraction; retry the upload: {diagnostic}" + ))), + Some(code) => Err(UploadError::Remote(format!( + "extraction exited with status {code}: {diagnostic}" + ))), } - let status = status.map_err(|error| UploadError::Transport(error.to_string()))?; - if !status.success() { - if status.code() == Some(super::SSH_TRANSPORT_FAILURE_EXIT_CODE) { - return Err(UploadError::Transport(format!("{status}: {diagnostic}"))); - } - return Err(UploadError::Protocol(format!( - "SSH exited with {status} after a successful result: {diagnostic}" - ))); - } - Ok(()) } async fn read_diagnostic(mut reader: R) -> String { @@ -180,45 +208,6 @@ async fn read_diagnostic(mut reader: R) -> String { String::from_utf8_lossy(&diagnostic).into_owned() } -fn protocol_error(error: ProtocolError) -> UploadError { - // ProtocolError preserves I/O errors separately from invalid frames. - match error { - ProtocolError::Io(error) => UploadError::Transport(error.to_string()), - error => UploadError::Protocol(error.to_string()), - } -} - -async fn negotiate( - reader: &mut R, - writer: &mut W, - plan: UploadPlan, -) -> Result<(u64, usize), UploadError> { - let hello = read_frame(reader).await.map_err(protocol_error)?; - let Frame::Hello { - window_bytes, - max_data_bytes, - } = hello - else { - return Err(server_failure(hello, "expected HELLO")); - }; - write_frame(writer, &Frame::Begin(plan)) - .await - .map_err(protocol_error)?; - match read_frame(reader).await.map_err(protocol_error)? { - Frame::Ready => Ok((u64::from(window_bytes), max_data_bytes as usize)), - frame => Err(server_failure(frame, "expected READY")), - } -} - -fn server_failure(frame: Frame, expected: &str) -> UploadError { - match frame { - Frame::Error(failure) => { - UploadError::Remote(format!("{:?}: {}", failure.code, failure.message)) - } - _ => UploadError::Protocol(expected.to_string()), - } -} - enum ProducerEvent { Data(Vec), Finished(Result<(), UploadError>), @@ -227,13 +216,12 @@ enum ProducerEvent { struct ChunkWriter { sender: mpsc::Sender, buffer: Vec, - max_chunk: usize, } impl ChunkWriter { fn flush_chunk(&mut self) -> io::Result<()> { if !self.buffer.is_empty() { - let buffer = std::mem::replace(&mut self.buffer, Vec::with_capacity(self.max_chunk)); + let buffer = std::mem::replace(&mut self.buffer, Vec::with_capacity(CHUNK_BYTES)); self.sender .blocking_send(ProducerEvent::Data(buffer)) .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "upload stopped"))?; @@ -246,10 +234,10 @@ impl Write for ChunkWriter { fn write(&mut self, mut bytes: &[u8]) -> io::Result { let original_len = bytes.len(); while !bytes.is_empty() { - let length = bytes.len().min(self.max_chunk - self.buffer.len()); + let length = bytes.len().min(CHUNK_BYTES - self.buffer.len()); self.buffer.extend_from_slice(&bytes[..length]); bytes = &bytes[length..]; - if self.buffer.len() == self.max_chunk { + if self.buffer.len() == CHUNK_BYTES { self.flush_chunk()?; } } @@ -261,7 +249,7 @@ impl Write for ChunkWriter { } } -fn spawn_producer

(producer: P, max_chunk: usize) -> mpsc::Receiver +fn spawn_producer

(producer: P) -> mpsc::Receiver where P: FnOnce(&mut dyn Write) -> miette::Result<()> + Send + 'static, { @@ -269,8 +257,7 @@ where tokio::task::spawn_blocking(move || { let mut writer = ChunkWriter { sender: sender.clone(), - buffer: Vec::with_capacity(max_chunk), - max_chunk, + buffer: Vec::with_capacity(CHUNK_BYTES), }; let result = producer(&mut writer) .map_err(|error| UploadError::Archive(error.to_string())) @@ -284,138 +271,117 @@ where receiver } -/// Dedicated readers and writers let remote failure interrupt a blocked local -/// SSH pipe write. The frame reader is never cancelled between partial reads. +/// Send the plan, wait for extraction to start, then stream the archive while +/// at most `WINDOW_BYTES` remain unacknowledged. Takes `stdin` only after the +/// archive is complete, so failures leave the caller to stop SSH first. async fn transfer( mut reader: R, - mut writer: W, - plan: UploadPlan, + stdin: &mut Option, + plan: Vec, producer: P, -) -> Result<(), UploadError> +) -> Result where R: AsyncRead + Unpin + Send + 'static, - W: AsyncWrite + Unpin + Send + 'static, + W: AsyncWrite + Unpin + Send, P: FnOnce(&mut dyn Write) -> miette::Result<()> + Send + 'static, { - let (window, max_chunk) = tokio::time::timeout( - NEGOTIATION_TIMEOUT, - negotiate(&mut reader, &mut writer, plan), - ) - .await - .map_err(|_| UploadError::Transport("upload negotiation timed out".to_string()))??; - let mut events = spawn_producer(producer, max_chunk); - let (incoming, mut frames) = mpsc::channel(4); - let reader_task = tokio::spawn(async move { - loop { - let frame = read_frame(&mut reader).await.map_err(protocol_error); - let terminal = frame.is_err(); - if incoming.send(frame).await.is_err() || terminal { - break; - } + let stream = |error: io::Error| Stopped::Stream(error.to_string()); + let writer = stdin.as_mut().expect("upload stdin is open"); + writer.write_all(&plan).await.map_err(stream)?; + writer.flush().await.map_err(stream)?; + match tokio::time::timeout(NEGOTIATION_TIMEOUT, reader.read_u64()).await { + Err(_) => { + return Err(Stopped::Local(UploadError::Transport( + "upload negotiation timed out".to_string(), + ))); } - }); - let (outgoing, mut queued) = mpsc::channel::(1); - let (completed, mut completions) = mpsc::channel(1); - let writer_task = tokio::spawn(async move { - while let Some(frame) = queued.recv().await { - let is_end = matches!(frame, Frame::End(_)); - let result = write_frame(&mut writer, &frame) - .await - .map_err(protocol_error); - let result = if result.is_ok() && is_end { - writer - .shutdown() - .await - .map_err(|error| UploadError::Transport(error.to_string())) - } else { - result - }; - let terminal = result.is_err() || is_end; - if completed.send(result).await.is_err() || terminal { - break; + Ok(Err(error)) => return Err(stream(error)), + Ok(Ok(0)) => {} + Ok(Ok(_)) => { + return Err(Stopped::Local(UploadError::Protocol( + "upload did not start at offset zero".to_string(), + ))); + } + } + // The reader is never cancelled between partial records. + let (totals, mut acknowledged) = watch::channel(0_u64); + let acks = tokio::spawn(async move { + while let Ok(total) = reader.read_u64().await { + if total < *totals.borrow() { + return Err(UploadError::Protocol( + "upload acknowledgment moved backwards".to_string(), + )); } + totals.send_replace(total); } + Ok(()) }); - let _tasks = AbortTasks(vec![reader_task.abort_handle(), writer_task.abort_handle()]); + let _acks = AbortTasks(vec![acks.abort_handle()]); + let mut chunks = spawn_producer(producer); let mut sent = 0_u64; - let mut acknowledged = 0_u64; - let mut digest = Sha256::new(); - let mut pending = None; - let mut producer_finished = false; - let mut writing = false; - let mut end = None; loop { - if !writing && end.is_none() { - if pending.as_ref().is_some_and(|data: &Vec| { - (sent - acknowledged) - .checked_add(data.len() as u64) - .is_some_and(|pending| pending <= window) - }) { - let data = pending.take().expect("pending upload data"); - sent = sent.checked_add(data.len() as u64).ok_or_else(|| { - UploadError::Protocol("upload byte count overflow".to_string()) - })?; - digest.update(&data); - outgoing - .send(Frame::Data(data)) - .await - .map_err(|_| UploadError::Transport("upload writer stopped".to_string()))?; - writing = true; - } else if producer_finished && pending.is_none() { - let summary = UploadSummary { - bytes: sent, - sha256: digest.clone().finalize().into(), - }; - outgoing - .send(Frame::End(summary)) - .await - .map_err(|_| UploadError::Transport("upload writer stopped".to_string()))?; - end = Some(summary); - writing = true; + let chunk = match chunks.recv().await { + Some(ProducerEvent::Data(chunk)) => chunk, + Some(ProducerEvent::Finished(result)) => { + result.map_err(Stopped::Local)?; + break; } - } - tokio::select! { - frame = frames.recv() => { - match frame.ok_or_else(|| UploadError::Transport("upload reader stopped".to_string()))?? { - Frame::Ack(bytes) => { - if bytes < acknowledged || bytes > sent { - return Err(UploadError::Protocol("invalid cumulative upload ACK".to_string())); - } - acknowledged = bytes; - } - Frame::Result(summary) => { - if end != Some(summary) { - return Err(UploadError::Protocol("upload result count or digest does not match END".to_string())); - } - return Ok(()); - } - frame => return Err(server_failure(frame, "unexpected frame during upload")), - } + None => { + return Err(Stopped::Local(UploadError::Archive( + "archive producer stopped".to_string(), + ))); } - result = completions.recv(), if writing => { - result.ok_or_else(|| UploadError::Transport("upload writer stopped".to_string()))??; - writing = false; + }; + let length = chunk.len() as u64; + loop { + let total = *acknowledged.borrow_and_update(); + if total > sent { + return Err(Stopped::Local(UploadError::Protocol( + "upload acknowledged unsent bytes".to_string(), + ))); } - event = events.recv(), if pending.is_none() && !producer_finished => { - match event.ok_or_else(|| UploadError::Archive("archive producer stopped".to_string()))? { - ProducerEvent::Data(data) => pending = Some(data), - ProducerEvent::Finished(result) => { - result?; - producer_finished = true; - } - } + if sent + length - total <= WINDOW_BYTES { + break; + } + if acknowledged.changed().await.is_err() { + return Err(match acks.await { + Ok(Err(error)) => Stopped::Local(error), + _ => Stopped::Stream("upload acknowledgments ended".to_string()), + }); } } + // Count bytes before writing: the supervisor may acknowledge part of + // a chunk while the pipe write is still in progress. + sent += length; + writer.write_all(&chunk).await.map_err(stream)?; } + // Closing stdin sends SSH EOF; the final total arrives before exit status. + writer.shutdown().await.map_err(stream)?; + stdin.take(); + if let Ok(Err(error)) = acks.await { + return Err(Stopped::Local(error)); + } + let acknowledged = *acknowledged.borrow(); + if acknowledged > sent { + return Err(Stopped::Local(UploadError::Protocol( + "upload acknowledged unsent bytes".to_string(), + ))); + } + Ok(Sent { + bytes: sent, + acknowledged, + }) } #[cfg(test)] mod tests { use super::*; - use openshell_core::upload::{FailureCode, MAX_DATA_BYTES, UploadFailure, WINDOW_BYTES}; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; - use tokio::io::{DuplexStream, ReadHalf, WriteHalf}; + use tokio::io::{AsyncBufReadExt, BufReader, DuplexStream, ReadHalf, WriteHalf}; + + const WINDOW: usize = 1024 * 1024; + type Check = fn(&Result<(), UploadError>) -> bool; fn plan() -> UploadPlan { UploadPlan { @@ -424,22 +390,30 @@ mod tests { } } - async fn accept(stream: DuplexStream) -> (ReadHalf, WriteHalf) { - let (mut read, mut write) = tokio::io::split(stream); - write_frame( - &mut write, - &Frame::Hello { - window_bytes: WINDOW_BYTES, - max_data_bytes: u32::try_from(MAX_DATA_BYTES).unwrap(), - }, - ) - .await - .unwrap(); - assert!(matches!( - read_frame(&mut read).await.unwrap(), - Frame::Begin(_) - )); - write_frame(&mut write, &Frame::Ready).await.unwrap(); + fn start( + producer: impl FnOnce(&mut dyn Write) -> miette::Result<()> + Send + 'static, + buffer: usize, + ) -> (tokio::task::JoinHandle>, DuplexStream) { + let (client, server) = tokio::io::duplex(buffer); + let (read, write) = tokio::io::split(client); + let plan = plan().encode().unwrap(); + let transfer = tokio::spawn(async move { + let mut write = Some(write); + transfer(read, &mut write, plan, producer).await + }); + (transfer, server) + } + + /// Read the plan line and grant the zero acknowledgment. + async fn accept( + server: DuplexStream, + ) -> (BufReader>, WriteHalf) { + let (read, mut write) = tokio::io::split(server); + let mut read = BufReader::new(read); + let mut line = Vec::new(); + read.read_until(b'\n', &mut line).await.unwrap(); + assert_eq!(line, plan().encode().unwrap()); + write.write_u64(0).await.unwrap(); (read, write) } @@ -449,62 +423,39 @@ mod tests { .collect() } - async fn complete( - read: &mut ReadHalf, - write: &mut WriteHalf, - mut received: Vec, - ) -> Vec { - loop { - match read_frame(read).await.unwrap() { - Frame::Data(data) => { - assert!(data.len() <= MAX_DATA_BYTES); - received.extend_from_slice(&data); - write_frame(write, &Frame::Ack(received.len() as u64)) - .await - .unwrap(); - } - Frame::End(summary) => { - assert_eq!(summary.bytes, received.len() as u64); - assert_eq!(summary.sha256, <[u8; 32]>::from(Sha256::digest(&received))); - write_frame(write, &Frame::Result(summary)).await.unwrap(); - return received; - } - frame => panic!("unexpected client frame: {frame:?}"), - } - } + fn produce(bytes: Vec) -> impl FnOnce(&mut dyn Write) -> miette::Result<()> { + move |writer| writer.write_all(&bytes).map_err(miette::Report::from_err) } #[tokio::test] async fn stalled_consumer_blocks_at_credit_limit_then_resumes_with_exact_data() { - let expected = payload(6 * WINDOW_BYTES as usize + 123); - let bytes = expected.clone(); - let (client, server) = tokio::io::duplex(4096); - let (read, write) = tokio::io::split(client); - let client = tokio::spawn(transfer(read, write, plan(), move |writer| { - writer.write_all(&bytes).map_err(miette::Report::from_err)?; - Ok(()) - })); + let expected = payload(6 * WINDOW + 123); + let (client, server) = start(produce(expected.clone()), 4096); let (mut read, mut write) = accept(server).await; - let mut received = Vec::new(); - while received.len() < WINDOW_BYTES as usize { - match read_frame(&mut read).await.unwrap() { - Frame::Data(data) => received.extend_from_slice(&data), - frame => panic!("unexpected client frame: {frame:?}"), - } - } - assert_eq!(received.len(), WINDOW_BYTES as usize); + let mut received = vec![0; WINDOW]; + read.read_exact(&mut received).await.unwrap(); assert!( - tokio::time::timeout(Duration::from_millis(50), read_frame(&mut read)) + tokio::time::timeout(Duration::from_millis(50), read.read_u8()) .await .is_err(), "client exceeded credit while the extractor was stalled" ); assert!(!client.is_finished()); - write_frame(&mut write, &Frame::Ack(received.len() as u64)) - .await - .unwrap(); - assert_eq!(complete(&mut read, &mut write, received).await, expected); - client.await.unwrap().unwrap(); + // Acknowledge in ordinary-sized steps until the client sends EOF. + let mut chunk = vec![0; 256 * 1024]; + loop { + write.write_u64(received.len() as u64).await.unwrap(); + let size = read.read(&mut chunk).await.unwrap(); + if size == 0 { + break; + } + received.extend_from_slice(&chunk[..size]); + } + write.shutdown().await.unwrap(); + let sent = client.await.unwrap().unwrap(); + assert_eq!(received, expected); + assert_eq!(sent.bytes, expected.len() as u64); + assert_eq!(sent.acknowledged, sent.bytes); } #[tokio::test] @@ -517,24 +468,18 @@ mod tests { } let finished = Arc::new(AtomicBool::new(false)); let producer_finished = finished.clone(); - let (client, server) = tokio::io::duplex(4096); - let (read, write) = tokio::io::split(client); - let client = tokio::spawn(transfer(read, write, plan(), move |writer| { - let _finished = ProducerExit(producer_finished); - writer - .write_all(&payload(8 * WINDOW_BYTES as usize)) - .map_err(miette::Report::from_err)?; - Ok(()) - })); + let (client, server) = start( + move |writer| { + let _finished = ProducerExit(producer_finished); + writer + .write_all(&payload(8 * WINDOW)) + .map_err(miette::Report::from_err) + }, + 4096, + ); let (mut read, _write) = accept(server).await; - let mut received = 0; - while received < WINDOW_BYTES as usize { - let Frame::Data(data) = read_frame(&mut read).await.unwrap() else { - panic!("expected data before credit is exhausted"); - }; - received += data.len(); - } - assert_eq!(received, WINDOW_BYTES as usize); + let mut received = vec![0; WINDOW]; + read.read_exact(&mut received).await.unwrap(); assert!(!finished.load(Ordering::SeqCst)); client.abort(); assert!(client.await.unwrap_err().is_cancelled()); @@ -556,150 +501,145 @@ mod tests { } #[tokio::test] - async fn impossible_ack_is_rejected_without_retry() { - let (client, server) = tokio::io::duplex(4096); - let (read, write) = tokio::io::split(client); - let client = tokio::spawn(transfer(read, write, plan(), |writer| { - writer - .write_all(&payload(2 * WINDOW_BYTES as usize)) - .map_err(miette::Report::from_err)?; - Ok(()) - })); - let (_read, mut write) = accept(server).await; - write_frame(&mut write, &Frame::Ack(u64::MAX)) - .await - .unwrap(); - let error = client.await.unwrap().unwrap_err(); - assert!(matches!(error, UploadError::Protocol(_))); - assert!(!error.is_retryable()); + async fn impossible_and_regressing_acks_are_rejected_without_retry() { + for regress in [false, true] { + let (client, server) = start(produce(payload(4 * WINDOW)), 4096); + let (mut read, mut write) = accept(server).await; + if regress { + let mut chunk = vec![0; CHUNK_BYTES]; + read.read_exact(&mut chunk).await.unwrap(); + write.write_u64(CHUNK_BYTES as u64).await.unwrap(); + write.write_u64(CHUNK_BYTES as u64 - 1).await.unwrap(); + } else { + write.write_u64(u64::MAX).await.unwrap(); + } + // Keep consuming so the client reaches its next credit decision. + let drain = tokio::spawn(async move { + let _ = tokio::io::copy(&mut read, &mut tokio::io::sink()).await; + }); + let Err(Stopped::Local(error)) = client.await.unwrap() else { + panic!("invalid acknowledgment was accepted"); + }; + assert!(matches!(error, UploadError::Protocol(_))); + assert!(!error.is_retryable()); + drop(write); + drain.await.unwrap(); + } } #[tokio::test] - async fn regressing_ack_is_rejected() { + async fn failed_archive_leaves_stdin_open_for_the_caller_to_stop_ssh() { let (client, server) = tokio::io::duplex(4096); let (read, write) = tokio::io::split(client); - let client = tokio::spawn(transfer(read, write, plan(), |writer| { - writer - .write_all(&payload(2 * WINDOW_BYTES as usize)) - .map_err(miette::Report::from_err)?; - Ok(()) - })); - let (mut read, mut write) = accept(server).await; - let Frame::Data(data) = read_frame(&mut read).await.unwrap() else { - panic!("expected data"); - }; - write_frame(&mut write, &Frame::Ack(data.len() as u64)) - .await - .unwrap(); - write_frame(&mut write, &Frame::Ack(data.len() as u64 - 1)) - .await - .unwrap(); - assert!(matches!( - client.await.unwrap().unwrap_err(), - UploadError::Protocol(_) - )); - } - - #[tokio::test] - async fn terminal_error_interrupts_blocked_pipe_write() { - let (client, server) = tokio::io::duplex(64); - let (read, write) = tokio::io::split(client); - let client = tokio::spawn(transfer(read, write, plan(), |writer| { + let mut stdin = Some(write); + let server = tokio::spawn(async move { + let (mut read, write) = accept(server).await; + let mut partial = [0; 7]; + read.read_exact(&mut partial).await.unwrap(); + (read, write) + }); + let result = transfer(read, &mut stdin, plan().encode().unwrap(), |writer| { writer - .write_all(&payload(2 * WINDOW_BYTES as usize)) + .write_all(b"partial") .map_err(miette::Report::from_err)?; - Ok(()) - })); - let (_read, mut write) = accept(server).await; - write_frame( - &mut write, - &Frame::Error(UploadFailure { - code: FailureCode::InputWrite, - message: "extractor closed stdin: connection is unavailable".to_string(), - }), - ) - .await - .unwrap(); - let error = tokio::time::timeout(Duration::from_secs(1), client) - .await - .expect("remote failure must interrupt a blocked writer") - .unwrap() - .unwrap_err(); - assert!(matches!(error, UploadError::Remote(_))); - assert!(!error.is_retryable()); - } - - #[tokio::test] - async fn successful_exit_cannot_replace_terminal_result() { - let (client, server) = tokio::io::duplex(4096); - let (read, write) = tokio::io::split(client); - let client = tokio::spawn(transfer(read, write, plan(), |_| Ok(()))); - let (mut read, write) = accept(server).await; + writer.flush().map_err(miette::Report::from_err)?; + Err(miette::miette!("source file vanished")) + }) + .await; assert!(matches!( - read_frame(&mut read).await.unwrap(), - Frame::End(_) + result, + Err(Stopped::Local(UploadError::Archive(_))) )); - drop(write); - drop(read); - assert!(client.await.unwrap().is_err()); + // EOF would mark a truncated archive as complete. + assert!(stdin.is_some()); + let (mut read, _write) = server.await.unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(50), read.read_u8()) + .await + .is_err() + ); } #[tokio::test] - async fn mismatched_terminal_digest_is_rejected() { - let (client, server) = tokio::io::duplex(4096); - let (read, write) = tokio::io::split(client); - let client = tokio::spawn(transfer(read, write, plan(), |_| Ok(()))); - let (mut read, mut write) = accept(server).await; - let Frame::End(mut summary) = read_frame(&mut read).await.unwrap() else { - panic!("expected end"); - }; - summary.sha256[0] ^= 1; - write_frame(&mut write, &Frame::Result(summary)) + async fn closed_stream_interrupts_blocked_pipe_write() { + let (client, server) = start(produce(payload(2 * WINDOW)), 64); + let (read, write) = accept(server).await; + drop((read, write)); + let result = tokio::time::timeout(Duration::from_secs(1), client) .await + .expect("remote closure must interrupt a blocked writer") .unwrap(); - assert!(matches!( - client.await.unwrap().unwrap_err(), - UploadError::Protocol(_) - )); + assert!(matches!(result, Err(Stopped::Stream(_)))); } #[cfg(unix)] - #[tokio::test] - async fn unavailable_subsystem_never_starts_archive_producer() { - let started = Arc::new(AtomicBool::new(false)); - let producer_started = started.clone(); + async fn run_script(script: &str, started: Arc) -> Result<(), UploadError> { let mut command = Command::new("sh"); - command.args([ - "-c", - "printf 'subsystem request failed on channel 0\\n' >&2; exit 255", - ]); - let error = run(command, plan(), move |_| { - producer_started.store(true, Ordering::SeqCst); - Ok(()) + command.args(["-c", script]); + run(command, plan(), move |writer| { + started.store(true, Ordering::SeqCst); + writer + .write_all(b"archive") + .map_err(miette::Report::from_err) }) .await - .unwrap_err(); - assert!(matches!(error, UploadError::Unsupported)); - assert!(!started.load(Ordering::SeqCst)); - assert!(!error.is_retryable()); } #[cfg(unix)] #[tokio::test] - async fn early_ssh_exit_never_starts_archive_producer() { - let started = Arc::new(AtomicBool::new(false)); - let producer_started = started.clone(); - let mut command = Command::new("sh"); - command.args(["-c", "exit 255"]); - let error = run(command, plan(), move |_| { - producer_started.store(true, Ordering::SeqCst); - Ok(()) - }) - .await - .unwrap_err(); - assert!(matches!(error, UploadError::Transport(_))); - assert!(!started.load(Ordering::SeqCst)); - assert!(error.is_retryable()); + async fn exit_status_and_final_total_decide_the_result() { + // Each script reads the plan line, grants zero credit, then drains the + // archive before reporting. `\0\0\0\0\0\0\0\a` is a total of seven bytes. + const ZERO: &str = r"printf '\0\0\0\0\0\0\0\0'"; + const SEVEN: &str = r"printf '\0\0\0\0\0\0\0\a'"; + assert_eq!(WINDOW as u64, WINDOW_BYTES); + let cases: [(String, Check); 4] = [ + ( + format!("read -r plan; {ZERO}; cat >/dev/null; {SEVEN}"), + |result| result.is_ok(), + ), + (format!("read -r plan; {ZERO}; cat >/dev/null"), |result| { + matches!(result, Err(UploadError::Protocol(_))) + }), + ( + format!("read -r plan; {ZERO}; cat >/dev/null; echo 'tar: boom' >&2; exit 2"), + |result| matches!(result, Err(UploadError::Remote(message)) if message.contains("tar: boom")), + ), + ( + format!("read -r plan; {ZERO}; cat >/dev/null; exit 73"), + |result| matches!(result, Err(UploadError::Remote(message)) if message.contains("retry the upload")), + ), + ]; + for (script, check) in cases { + let started = Arc::new(AtomicBool::new(false)); + let result = run_script(&script, started.clone()).await; + assert!(check(&result), "{script}: {result:?}"); + assert!(started.load(Ordering::SeqCst)); + if let Err(error) = result { + assert!(!error.is_retryable()); + } + } + } + + #[cfg(unix)] + #[tokio::test] + async fn failure_before_extraction_starts_never_starts_archive_producer() { + for (script, retryable) in [ + ( + "printf 'subsystem request failed on channel 0\\n' >&2; exit 255", + false, + ), + ("exit 255", true), + ("read -r plan; echo 'mkdir: denied' >&2; exit 1", false), + ] { + let started = Arc::new(AtomicBool::new(false)); + let error = run_script(script, started.clone()).await.unwrap_err(); + assert!(!started.load(Ordering::SeqCst), "{script}"); + assert_eq!(error.is_retryable(), retryable, "{script}: {error:?}"); + if script.contains("subsystem") { + assert!(matches!(error, UploadError::Unsupported)); + } + } } #[cfg(unix)] diff --git a/crates/openshell-core/src/upload.rs b/crates/openshell-core/src/upload.rs index 588159e0f0..f3939a5aa4 100644 --- a/crates/openshell-core/src/upload.rs +++ b/crates/openshell-core/src/upload.rs @@ -3,33 +3,29 @@ //! Internal streaming-upload protocol shared by the CLI and SSH supervisor. //! -//! SSH receive credit is replenished before application consumption. This -//! subsystem supplies its own bounded credit: an ACK means boundary stdin has -//! accepted bytes, while RESULT means extraction completed successfully. +//! SSH receive credit is replenished before application consumption, so this +//! subsystem returns its own credit and relies on SSH for everything else. The +//! client sends one JSON plan line, the raw archive, then EOF. The supervisor's +//! stdout carries only big-endian `u64` totals of archive bytes written to +//! extraction stdin; the first total, zero, means extraction started. +//! Extraction output arrives on stderr and the SSH exit status is the result. use serde::{Deserialize, Serialize}; -use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use crate::forward::shell_escape; pub const SUBSYSTEM: &str = "openshell-upload-v1"; -pub const WINDOW_BYTES: u32 = 1024 * 1024; -pub const MAX_DATA_BYTES: usize = 64 * 1024; -pub const MAX_METADATA_BYTES: usize = 16 * 1024; -pub const MAX_DIAGNOSTIC_BYTES: usize = 16 * 1024; +/// Archive bytes a client may send beyond the last acknowledged total. This +/// must stay below the supervisor's pending-stdin budget. +pub const WINDOW_BYTES: u64 = 1024 * 1024; +pub const MAX_PLAN_BYTES: usize = 16 * 1024; pub const DESTINATION_CHANGED_EXIT_CODE: i32 = 73; -const HEADER_BYTES: usize = 12; -const MAGIC: &[u8; 4] = b"OSUP"; -// JSON can expand one control character into six bytes (a Unicode escape). -const MAX_ERROR_BYTES: usize = MAX_DIAGNOSTIC_BYTES * 6 + 512; #[derive(Debug, thiserror::Error)] -pub enum ProtocolError { - #[error("upload transport: {0}")] - Io(#[from] std::io::Error), - #[error("invalid upload protocol: {0}")] +pub enum PlanError { + #[error("invalid upload plan: {0}")] Invalid(&'static str), - #[error("invalid upload metadata: {0}")] + #[error("invalid upload plan: {0}")] Json(#[from] serde_json::Error), } @@ -67,252 +63,59 @@ pub struct UploadPlan { } impl UploadPlan { - pub fn validate(&self) -> Result<(), ProtocolError> { + pub fn validate(&self) -> Result<(), PlanError> { validate_path(&self.destination_directory)?; if let Some(expected) = &self.expected_destination { validate_path(&expected.path)?; } Ok(()) } -} - -fn validate_path(path: &str) -> Result<(), ProtocolError> { - if path.is_empty() || path.len() > MAX_METADATA_BYTES || path.contains('\0') { - return Err(ProtocolError::Invalid("invalid destination path")); - } - Ok(()) -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub struct UploadSummary { - pub bytes: u64, - pub sha256: [u8; 32], -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] -#[serde(rename_all = "snake_case")] -pub enum FailureCode { - Protocol, - DestinationChanged, - Extract, - InputWrite, - Internal, -} - -#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct UploadFailure { - pub code: FailureCode, - pub message: String, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum Frame { - Hello { - window_bytes: u32, - max_data_bytes: u32, - }, - Begin(UploadPlan), - Ready, - Data(Vec), - Ack(u64), - End(UploadSummary), - Result(UploadSummary), - Error(UploadFailure), - Cancel, -} -impl Frame { - pub fn encode(&self) -> Result, ProtocolError> { - let (kind, payload) = match self { - Self::Hello { - window_bytes, - max_data_bytes, - } => { - validate_limits(*window_bytes, *max_data_bytes)?; - let mut bytes = window_bytes.to_be_bytes().to_vec(); - bytes.extend_from_slice(&max_data_bytes.to_be_bytes()); - (1, bytes) - } - Self::Begin(plan) => { - plan.validate()?; - (2, serde_json::to_vec(plan)?) - } - Self::Ready => (3, Vec::new()), - Self::Data(bytes) => (4, bytes.clone()), - Self::Ack(bytes) => (5, bytes.to_be_bytes().to_vec()), - Self::End(summary) => (6, encode_summary(*summary)), - Self::Result(summary) => (7, encode_summary(*summary)), - Self::Error(failure) => { - if failure.message.len() > MAX_DIAGNOSTIC_BYTES { - return Err(ProtocolError::Invalid("diagnostic exceeds limit")); - } - (8, serde_json::to_vec(failure)?) - } - Self::Cancel => (9, Vec::new()), - }; - validate_length(kind, payload.len())?; - let length = u32::try_from(payload.len()) - .map_err(|_| ProtocolError::Invalid("frame exceeds limit"))?; - let mut encoded = Vec::with_capacity(HEADER_BYTES + payload.len()); - encoded.extend_from_slice(MAGIC); - encoded.extend_from_slice(&[kind, 0, 0, 0]); - encoded.extend_from_slice(&length.to_be_bytes()); - encoded.extend_from_slice(&payload); - Ok(encoded) + /// Encode the newline-terminated plan that precedes the archive. JSON + /// escapes newlines inside strings, so the first newline ends the plan. + pub fn encode(&self) -> Result, PlanError> { + self.validate()?; + let mut line = serde_json::to_vec(self)?; + if line.len() > MAX_PLAN_BYTES { + return Err(PlanError::Invalid("plan exceeds limit")); + } + line.push(b'\n'); + Ok(line) } } -fn validate_limits(window: u32, chunk: u32) -> Result<(), ProtocolError> { - if window == 0 - || window > WINDOW_BYTES - || chunk == 0 - || chunk as usize > MAX_DATA_BYTES - || chunk > window - { - return Err(ProtocolError::Invalid("invalid credit limits")); +fn validate_path(path: &str) -> Result<(), PlanError> { + if path.is_empty() || path.contains('\0') { + return Err(PlanError::Invalid("invalid destination path")); } Ok(()) } -fn encode_summary(summary: UploadSummary) -> Vec { - let mut encoded = summary.bytes.to_be_bytes().to_vec(); - encoded.extend_from_slice(&summary.sha256); - encoded -} - -fn validate_length(kind: u8, length: usize) -> Result<(), ProtocolError> { - let valid = match kind { - 1 | 5 => length == 8, - 2 => length > 0 && length <= MAX_METADATA_BYTES, - 3 | 9 => length == 0, - 4 => length > 0 && length <= MAX_DATA_BYTES, - 6 | 7 => length == 40, - 8 => length > 0 && length <= MAX_ERROR_BYTES, - _ => return Err(ProtocolError::Invalid("unknown frame type")), - }; - if !valid { - return Err(ProtocolError::Invalid("invalid frame length")); - } - Ok(()) -} - -fn parse_header(header: &[u8; HEADER_BYTES]) -> Result<(u8, usize), ProtocolError> { - if &header[..4] != MAGIC || header[5..8] != [0, 0, 0] { - return Err(ProtocolError::Invalid("invalid frame header")); - } - let kind = header[4]; - let length = u32::from_be_bytes(header[8..12].try_into().expect("fixed length")) as usize; - validate_length(kind, length)?; - Ok((kind, length)) -} - -fn decode_payload(kind: u8, payload: Vec) -> Result { - Ok(match kind { - 1 => { - let window_bytes = u32::from_be_bytes(payload[..4].try_into().expect("fixed length")); - let max_data_bytes = - u32::from_be_bytes(payload[4..8].try_into().expect("fixed length")); - validate_limits(window_bytes, max_data_bytes)?; - Frame::Hello { - window_bytes, - max_data_bytes, - } - } - 2 => { - let plan: UploadPlan = serde_json::from_slice(&payload)?; - plan.validate()?; - Frame::Begin(plan) - } - 3 => Frame::Ready, - 4 => Frame::Data(payload), - 5 => Frame::Ack(u64::from_be_bytes( - payload[..8].try_into().expect("fixed length"), - )), - 6 | 7 => { - let summary = UploadSummary { - bytes: u64::from_be_bytes(payload[..8].try_into().expect("fixed length")), - sha256: payload[8..40].try_into().expect("fixed length"), - }; - if kind == 6 { - Frame::End(summary) - } else { - Frame::Result(summary) - } - } - 8 => { - let failure: UploadFailure = serde_json::from_slice(&payload)?; - if failure.message.len() > MAX_DIAGNOSTIC_BYTES { - return Err(ProtocolError::Invalid("diagnostic exceeds limit")); - } - Frame::Error(failure) - } - 9 => Frame::Cancel, - _ => return Err(ProtocolError::Invalid("unknown frame type")), - }) -} - -/// Incremental decoder holding at most one bounded payload. The caller keeps -/// processing the remaining slice, so coalesced packets do not grow a queue. +/// Accumulates the plan line across SSH packets without exceeding the limit. #[derive(Default)] -pub struct Decoder { - header: [u8; HEADER_BYTES], - header_used: usize, - payload: Vec, - expected: Option<(u8, usize)>, -} - -impl Decoder { - pub fn decode(&mut self, input: &mut &[u8]) -> Result, ProtocolError> { - if self.header_used < HEADER_BYTES { - let size = input.len().min(HEADER_BYTES - self.header_used); - self.header[self.header_used..self.header_used + size].copy_from_slice(&input[..size]); - self.header_used += size; - *input = &input[size..]; - if self.header_used < HEADER_BYTES { - return Ok(None); - } - let expected = parse_header(&self.header).inspect_err(|_| { - self.header_used = 0; - })?; - self.payload.reserve(expected.1); - self.expected = Some(expected); +pub struct PlanReader(Vec); + +impl PlanReader { + /// Consume plan bytes from `input`, leaving any archive bytes after the + /// newline in place. Returns the plan once its line is complete. + pub fn read(&mut self, input: &mut &[u8]) -> Result, PlanError> { + let end = input.iter().position(|byte| *byte == b'\n'); + let line = &input[..end.unwrap_or(input.len())]; + if self.0.len() + line.len() > MAX_PLAN_BYTES { + return Err(PlanError::Invalid("plan exceeds limit")); } - let (kind, length) = self.expected.expect("validated header"); - let size = input.len().min(length - self.payload.len()); - self.payload.extend_from_slice(&input[..size]); - *input = &input[size..]; - if self.payload.len() < length { + self.0.extend_from_slice(line); + let Some(end) = end else { + *input = &[]; return Ok(None); - } - self.header_used = 0; - self.expected = None; - decode_payload(kind, std::mem::take(&mut self.payload)).map(Some) - } - - #[must_use] - pub const fn is_empty(&self) -> bool { - self.header_used == 0 + }; + *input = &input[end + 1..]; + let plan: UploadPlan = serde_json::from_slice(&self.0)?; + plan.validate()?; + Ok(Some(plan)) } } -pub async fn read_frame(reader: &mut R) -> Result { - let mut header = [0_u8; HEADER_BYTES]; - reader.read_exact(&mut header).await?; - let (kind, length) = parse_header(&header)?; - let mut payload = vec![0; length]; - reader.read_exact(&mut payload).await?; - decode_payload(kind, payload) -} - -pub async fn write_frame( - writer: &mut W, - frame: &Frame, -) -> Result<(), ProtocolError> { - writer.write_all(&frame.encode()?).await?; - Ok(()) -} - // This is the existing CLI destination probe, now shared with the receiving // supervisor. Check existing parents before treating failed existence as absence. const DESTINATION_KIND_SH: &str = r#"upload_destination_kind() { @@ -361,7 +164,7 @@ pub fn build_destination_probe(path: &str) -> String { format!("{DESTINATION_KIND_SH}\n{}", destination_invocation(path)) } -pub fn build_extract_command(plan: &UploadPlan) -> Result { +pub fn build_extract_command(plan: &UploadPlan) -> Result { plan.validate()?; let dest = shell_escape(&plan.destination_directory); let extract = format!("mkdir -p -- {dest} && cat | tar xf - -C {dest}"); @@ -387,116 +190,56 @@ pub fn build_extract_command(plan: &UploadPlan) -> Result mod tests { use super::*; - fn frames() -> Vec { - let summary = UploadSummary { - bytes: u64::MAX, - sha256: [7; 32], - }; - vec![ - Frame::Hello { - window_bytes: WINDOW_BYTES, - max_data_bytes: u32::try_from(MAX_DATA_BYTES).unwrap(), - }, - Frame::Begin(UploadPlan { - destination_directory: "/tmp/with 'quotes'".into(), - expected_destination: Some(ExpectedDestination { - path: "/tmp/file".into(), - kind: DestinationKind::Missing, - }), + fn plan() -> UploadPlan { + UploadPlan { + destination_directory: "/tmp/with 'quotes'\nand newline".into(), + expected_destination: Some(ExpectedDestination { + path: "/tmp/file".into(), + kind: DestinationKind::Missing, }), - Frame::Ready, - Frame::Data(vec![42; MAX_DATA_BYTES]), - Frame::Ack(u64::MAX), - Frame::End(summary), - Frame::Result(summary), - Frame::Error(UploadFailure { - code: FailureCode::InputWrite, - message: "stdin closed".into(), - }), - Frame::Cancel, - ] + } } #[test] - fn fragmented_and_coalesced_frames_round_trip() { - let frames = frames(); - let encoded = frames - .iter() - .flat_map(|f| f.encode().unwrap()) - .collect::>(); - for size in [1, 7, 12, 4096, encoded.len()] { - let mut decoder = Decoder::default(); - let mut output = Vec::new(); - for chunk in encoded.chunks(size) { + fn fragmented_plan_line_leaves_archive_bytes() { + let mut stream = plan().encode().unwrap(); + stream.extend_from_slice(b"archive\nbytes"); + for size in [1, 7, stream.len()] { + let mut reader = PlanReader::default(); + let mut decoded = None; + let mut rest = Vec::new(); + for chunk in stream.chunks(size) { let mut input = chunk; - while !input.is_empty() { - if let Some(frame) = decoder.decode(&mut input).unwrap() { - output.push(frame); - } + if decoded.is_none() { + decoded = reader.read(&mut input).unwrap(); } + rest.extend_from_slice(input); } - assert!(decoder.is_empty()); - assert_eq!(output, frames); + assert_eq!(decoded, Some(plan())); + assert_eq!(rest, b"archive\nbytes"); } } #[test] - fn malformed_headers_fail_before_payload_allocation() { - let good = Frame::Ready.encode().unwrap(); - for offset in [0, 4, 5, 6, 7, 8, 9, 10, 11] { - let mut bad = good.clone(); - bad[offset] = 255; - let mut decoder = Decoder::default(); - assert!(decoder.decode(&mut bad.as_slice()).is_err()); - assert_eq!(decoder.payload.capacity(), 0); - } - assert!(Frame::Data(vec![0; MAX_DATA_BYTES + 1]).encode().is_err()); - assert!( - Frame::Hello { - window_bytes: WINDOW_BYTES + 1, - max_data_bytes: 1 - } - .encode() - .is_err() - ); - } - - #[test] - fn metadata_rejects_unknown_fields_and_nul_paths() { - let bad = br#"{"destination_directory":"/tmp","expected_destination":null,"command":"id"}"# - .to_vec(); - assert!(decode_payload(2, bad).is_err()); - assert!( - UploadPlan { - destination_directory: "/tmp\0bad".into(), - expected_destination: None - } - .validate() - .is_err() - ); - assert!( - Frame::Error(UploadFailure { - code: FailureCode::Protocol, - message: "a".repeat(MAX_DIAGNOSTIC_BYTES + 1) - }) - .encode() - .is_err() - ); - } - - #[tokio::test] - async fn async_codec_rejects_truncation() { - let encoded = Frame::Data(vec![7; 32]).encode().unwrap(); - let mut input = &encoded[..encoded.len() - 1]; - assert!(matches!( - read_frame(&mut input).await, - Err(ProtocolError::Io(_)) - )); - let mut output = Vec::new(); - write_frame(&mut output, &Frame::Ack(12)).await.unwrap(); - assert_eq!( - read_frame(&mut output.as_slice()).await.unwrap(), - Frame::Ack(12) - ); + fn plan_rejects_unknown_fields_nul_paths_and_unbounded_lines() { + let mut input = + &br#"{"destination_directory":"/tmp","expected_destination":null,"command":"id"} +"#[..]; + assert!(PlanReader::default().read(&mut input).is_err()); + let nul = UploadPlan { + destination_directory: "/tmp\0bad".into(), + expected_destination: None, + }; + assert!(nul.encode().is_err()); + let long = UploadPlan { + destination_directory: "/".repeat(MAX_PLAN_BYTES), + expected_destination: None, + }; + assert!(long.encode().is_err()); + let unterminated = vec![b'{'; MAX_PLAN_BYTES + 1]; + let mut reader = PlanReader::default(); + let (first, second) = unterminated.split_at(MAX_PLAN_BYTES); + assert!(reader.read(&mut &first[..]).unwrap().is_none()); + assert!(reader.read(&mut &second[..]).is_err()); } } diff --git a/crates/openshell-supervisor-process/README.md b/crates/openshell-supervisor-process/README.md index 7d6e6f13db..7f7b43040a 100644 --- a/crates/openshell-supervisor-process/README.md +++ b/crates/openshell-supervisor-process/README.md @@ -5,26 +5,32 @@ isolation backend. ## Bounded streaming uploads -The CLI's built-in uploads use the `openshell-upload-v1` SSH subsystem. The -gateway relays its SSH channel without interpreting upload frames. The -supervisor launches the extractor through the ordinary `BoundaryExec` path, -with the same workload identity, environment, working directory, and -filesystem policy as other SSH children. +The CLI's built-in uploads use the `openshell-upload-v1` SSH subsystem. Russh +returns SSH receive credit before the supervisor consumes channel data, so this +subsystem returns its own credit and uses SSH for everything else. The gateway +relays the channel without interpreting it. -The upload protocol reserves stdout for control frames and the final result. -Extractor output is drained into bounded diagnostics. The supervisor grants a -1 MiB application window, and data frames carry at most 64 KiB. The CLI stops -sending archive data when outstanding bytes reach the window. A cumulative -acknowledgment returns credit only after `InputTask` successfully writes those -bytes through `BoundaryInput::write_all` and clears them from its in-flight -budget. The existing 4 MiB pending-input guard remains active. +The client sends one JSON plan line with the destination, then the raw tar +stream, then SSH EOF. When the plan arrives, the supervisor starts the +extractor through the ordinary `BoundaryExec` path, with the same workload +identity, environment, working directory and filesystem policy as other SSH +children. Startup runs inline, like an exec request. -The end frame carries the archive byte count and SHA-256 digest. A successful -result requires a matching stream, drained accepted writes, and a successful -extractor exit. Protocol failure, extractor failure, or channel cancellation -terminates the upload process and stops granting credit. Extraction writes -directly to the destination, so failures can leave partially extracted files; -the protocol does not provide rollback or atomic replacement. +Stdout carries only big-endian `u64` totals of archive bytes that `InputTask` +has written through `BoundaryInput::write_all`. The first total, zero, means +extraction started. Extractor stdout and stderr are both sent as SSH stderr, +so extractor output cannot forge credit. The client keeps at most 1 MiB beyond +the last total outstanding; the supervisor rejects input beyond that window. +The existing 4 MiB pending-input guard remains active. + +After the extractor exits successfully and stdin has drained, the supervisor +sends the final total, then the exit status. The client succeeds only when the +status is zero and the final total equals the bytes it sent. A malformed plan, +credit violation, extended data or EOF before the plan fails the upload with +exit status 74. Whichever happens first, the exit or the rejection, decides the +reported result. Rejection, channel close or disconnect terminates the +extractor. Extraction writes directly to the destination, so failures can +leave partially extracted files; there is no rollback or atomic replacement. The CLI requires supervisor support for this subsystem before sending archive data. It does not fall back to the previous unpaced exec-stdin upload path. diff --git a/crates/openshell-supervisor-process/src/ssh.rs b/crates/openshell-supervisor-process/src/ssh.rs index 18fd521ab3..7e2601c34f 100644 --- a/crates/openshell-supervisor-process/src/ssh.rs +++ b/crates/openshell-supervisor-process/src/ssh.rs @@ -375,7 +375,7 @@ async fn handle_connection( #[allow(clippy::struct_excessive_bools)] #[derive(Default)] struct ChannelState { - upload: Option, + upload: Option, command_started: bool, input_sender: Option, input_task: Option, @@ -737,12 +737,8 @@ impl russh::server::Handler for SshHandler { session.channel_failure(channel)?; return Ok(()); } - state.upload = Some(upload::UploadChannel::spawn( - channel, - session.handle(), - self.boundary_exec.clone(), - !state.no_login_shell, - )); + // Extraction starts once the plan line arrives on stdin. + state.upload = Some(upload::Upload::new()); session.channel_success(channel)?; return Ok(()); } @@ -929,13 +925,12 @@ impl russh::server::Handler for SshHandler { data: &[u8], session: &mut Session, ) -> Result<(), Self::Error> { - if let Some(upload) = self + if self .channels - .get_mut(&channel) - .and_then(|state| state.upload.as_mut()) + .get(&channel) + .is_some_and(|state| state.upload.is_some()) { - upload.data(data); - return Ok(()); + return self.upload_data(channel, data, session).await; } let Some(state) = self.channels.get_mut(&channel) else { warn!("data on unknown channel {channel:?}"); @@ -948,7 +943,7 @@ impl russh::server::Handler for SshHandler { if let Some(InputSender::Process(sender)) = &state.input_sender && sender.send(data) == Err(input::SendError::Full) { - self.reject_exec_input(channel, session)?; + self.reject_exec_input(channel, session, "process stdin buffer is full")?; } return Ok(()); } @@ -1001,16 +996,20 @@ impl russh::server::Handler for SshHandler { async fn channel_eof( &mut self, channel: ChannelId, - _session: &mut Session, + session: &mut Session, ) -> Result<(), Self::Error> { + if self + .channels + .get(&channel) + .and_then(|state| state.upload.as_ref()) + .is_some_and(upload::Upload::awaiting_plan) + { + return self.reject_upload(channel, session, Some("upload ended before its plan")); + } // Drop only the sender: the writer drains accepted bytes and then // closes stdin. Commands such as `cat | tar xf -` need this half-close // to finish while their output remains available to the SSH client. if let Some(state) = self.channels.get_mut(&channel) { - if let Some(upload) = state.upload.as_mut() { - upload.eof(); - return Ok(()); - } if state.main_attached && let Some(owner) = state.main_input_owner.take() && let Some(main_session) = self.main_session.as_ref() @@ -1034,14 +1033,14 @@ impl russh::server::Handler for SshHandler { channel: ChannelId, _ext: u32, _data: &[u8], - _session: &mut Session, + session: &mut Session, ) -> Result<(), Self::Error> { - if let Some(upload) = self + if self .channels - .get_mut(&channel) - .and_then(|state| state.upload.as_mut()) + .get(&channel) + .is_some_and(|state| state.upload.is_some()) { - upload.reject_extended_data(); + return self.reject_upload(channel, session, None); } Ok(()) } @@ -1072,12 +1071,11 @@ impl russh::server::Handler for SshHandler { } return Ok(()); } - let Some(process) = self.channels.get(&channel).and_then(|state| { - state - .process - .clone() - .or_else(|| state.upload.as_ref()?.process()) - }) else { + let Some(process) = self + .channels + .get(&channel) + .and_then(|state| state.process.clone()) + else { return Ok(()); }; let signal = match signal { @@ -1097,10 +1095,100 @@ impl russh::server::Handler for SshHandler { } impl SshHandler { + /// Read the plan, start extraction, then admit archive bytes within credit. + async fn upload_data( + &mut self, + channel: ChannelId, + mut data: &[u8], + session: &mut Session, + ) -> anyhow::Result<()> { + let Some(state) = self.channels.get_mut(&channel) else { + return Ok(()); + }; + let Some(upload) = state.upload.as_mut() else { + return Ok(()); + }; + if upload.awaiting_plan() { + let command = match upload.read_plan(&mut data).and_then(|plan| { + plan.map(|plan| openshell_core::upload::build_extract_command(&plan)) + .transpose() + }) { + Ok(None) => return Ok(()), + Ok(Some(command)) => command, + Err(error) => { + return self.reject_upload(channel, session, Some(&error.to_string())); + } + }; + let spec = openshell_isolation_interface::contract::ExecSpec { + program: String::new(), + args: Vec::new(), + shell: Some(openshell_isolation_interface::contract::ShellSpec { + command: Some(command), + login: !state.no_login_shell, + }), + runtime_helper: None, + env: Vec::new(), + workdir: None, + pty: false, + }; + // Like exec requests, start inline: a close is handled only after + // the process is registered with this channel. + if let Err(error) = self.start_exec_spec(channel, session.handle(), spec).await { + let message = format!("failed to start upload extraction: {error}"); + return self.reject_upload(channel, session, Some(&message)); + } + if data.is_empty() { + return Ok(()); + } + } + let Some(state) = self.channels.get_mut(&channel) else { + return Ok(()); + }; + let (Some(upload), Some(InputSender::Process(sender))) = + (state.upload.as_mut(), state.input_sender.as_ref()) + else { + return Ok(()); + }; + if !upload.admits(data.len()) { + return self.reject_upload(channel, session, None); + } + match sender.send(data) { + Ok(()) => upload.record(data.len()), + Err(input::SendError::Full) => return self.reject_upload(channel, session, None), + // The extractor closed stdin; its exit status reports the failure. + Err(input::SendError::Closed) => {} + } + Ok(()) + } + + /// Fail an upload with exit status 74 unless its result was reported. + fn reject_upload( + &mut self, + channel: ChannelId, + session: &mut Session, + message: Option<&str>, + ) -> anyhow::Result<()> { + if !self + .channels + .get(&channel) + .and_then(|state| state.upload.as_ref()) + .is_some_and(upload::Upload::reject) + { + return Ok(()); + } + // Only startup failures carry a message; the client's output window is + // still open then. Later rejections must not queue output. + if let Some(message) = message { + session.extended_data(channel, 1, format!("openshell: {message}\n").into_bytes())?; + } + self.reject_exec_input(channel, session, "invalid upload stream") + } + fn reject_exec_input( &mut self, channel: ChannelId, session: &mut Session, + reason: &str, ) -> anyhow::Result<()> { if let Some(mut state) = self.channels.remove(&channel) { // A locally initiated close may never call channel_close. Release @@ -1117,7 +1205,7 @@ impl SshHandler { // Session methods enqueue directly. Do not add stderr data here: with // no output credit, russh would retain it and stop draining Handle // messages for other channels. Exit status needs no output credit. - warn!(?channel, "process stdin buffer is full; terminating exec"); + warn!(?channel, "{reason}; terminating exec"); session.exit_status_request(channel, 74)?; session.eof(channel)?; session.close(channel)?; @@ -1189,13 +1277,28 @@ impl SshHandler { .channels .get_mut(&channel) .ok_or_else(|| anyhow::anyhow!("exec on unknown channel {channel:?}"))?; + if state.upload.is_some() && exec.stdin.is_none() { + let process = exec.process.clone(); + tokio::spawn(async move { + let _ = process.terminate().await; + }); + anyhow::bail!("extraction process has no stdin"); + } state.process = Some(exec.process.clone()); state.command_started = true; state.terminal = exec.terminal.take(); let output_status = exec.output_status.take(); + let mut finisher = None; if let Some(stdin) = exec.stdin.take() { - let (sender, task) = input::InputTask::spawn(stdin); + let (sender, task) = if let Some(upload) = state.upload.as_mut() { + let (sender, task, progress) = input::InputTask::spawn_with_progress(stdin); + finisher = + Some(upload.start(channel, handle.clone(), progress, exec.process.clone())); + (sender, task) + } else { + input::InputTask::spawn(stdin) + }; state.input_sender = Some(InputSender::Process(sender)); state.input_task = Some(task); } @@ -1204,6 +1307,8 @@ impl SshHandler { .as_ref() .map(input::InputTask::abort_handle); + // Upload stdout carries only credit, so extraction output uses stderr. + let stdout_extended = finisher.is_some(); let mut stdout = exec.stdout; let stdout_handle = handle.clone(); let stdout_task = tokio::spawn(async move { @@ -1212,7 +1317,12 @@ impl SshHandler { match stdout.read(&mut buffer).await { Ok(0) | Err(_) => break, Ok(size) => { - let _ = stdout_handle.data(channel, buffer[..size].to_vec()).await; + let data = buffer[..size].to_vec(); + let _ = if stdout_extended { + stdout_handle.extended_data(channel, 1, data).await + } else { + stdout_handle.data(channel, data).await + }; } } } @@ -1250,9 +1360,6 @@ impl SshHandler { } status }; - if let Some(input_abort) = input_abort { - input_abort.abort(); - } let code = match status { Some(openshell_isolation_interface::contract::BoundaryExitStatus::Exited(code)) => { code.max(0).cast_unsigned() @@ -1262,6 +1369,16 @@ impl SshHandler { )) => (128_i32.saturating_add(signal)).max(0).cast_unsigned(), None => 74, }; + let code = match finisher { + Some(finisher) => finisher.finish(&handle, code).await, + None => Some(code), + }; + if let Some(input_abort) = input_abort { + input_abort.abort(); + } + let Some(code) = code else { + return; + }; let _ = handle.eof(channel).await; let _ = handle.exit_status_request(channel, code).await; let _ = handle.close(channel).await; diff --git a/crates/openshell-supervisor-process/src/ssh/upload.rs b/crates/openshell-supervisor-process/src/ssh/upload.rs index 21fbc532c7..11b3363719 100644 --- a/crates/openshell-supervisor-process/src/ssh/upload.rs +++ b/crates/openshell-supervisor-process/src/ssh/upload.rs @@ -1,240 +1,120 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Upload input credit follows boundary consumption, independently of russh's -//! receive window. The SSH callback only decodes and queues bounded input. - -use super::input::{InputProgress, InputSender, InputState, InputTask, SendError}; -use openshell_core::upload::{ - DESTINATION_CHANGED_EXIT_CODE, Decoder, FailureCode, Frame, MAX_DATA_BYTES, - MAX_DIAGNOSTIC_BYTES, UploadFailure, UploadPlan, UploadSummary, WINDOW_BYTES, - build_extract_command, -}; -use openshell_isolation_interface::contract::{ - BoundaryExec, BoundaryExitStatus, BoundaryProcess, ExecSession, ExecSpec, ShellSpec, -}; +//! Upload channels are exec channels whose stdin credit follows consumption. +//! +//! Russh replenishes SSH receive credit before `data()` runs, so the channel +//! reports how many archive bytes extraction stdin accepted. Stdout carries +//! only those totals; extraction output is sent as stderr so it cannot forge +//! credit. The SSH exit status reports the result. + +use super::input::{InputProgress, InputState}; +use openshell_core::upload::{PlanError, PlanReader, UploadPlan, WINDOW_BYTES}; +use openshell_isolation_interface::contract::BoundaryProcess; use russh::ChannelId; use russh::server::Handle; -use sha2::{Digest, Sha256}; -use std::sync::{Arc, Mutex}; -use tokio::io::{AsyncRead, AsyncReadExt}; -use tokio::sync::{Notify, oneshot, watch}; -use tokio::task::JoinHandle; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; +use tokio::sync::watch; +use tokio::task::{AbortHandle, JoinHandle}; const ACK_INTERVAL: u64 = 256 * 1024; +/// A successful `cat | tar` exit follows stdin EOF, so the writer publishes +/// its drained state almost immediately. +const DRAIN_TIMEOUT: Duration = Duration::from_secs(5); -#[derive(Clone, Copy, PartialEq, Eq)] -enum Phase { - AwaitBegin, - Starting, - Receiving, - Draining, - Finished, -} - -struct State { - phase: Phase, - plan: Option, - sender: Option, - input_task: Option, +pub(super) struct Upload { + plan: Option, + received: u64, progress: Option>, process: Option>, - received: u64, - digest: Sha256, - summary: Option, - failure: Option, + acks: Option, + /// Set by whichever side reports the result: the exit task or a rejection. + reported: Arc, } -struct Shared { - state: Mutex, - changed: Notify, -} - -pub(super) struct UploadChannel { - decoder: Decoder, - shared: Arc, - owner: JoinHandle<()>, -} - -impl UploadChannel { - pub(super) fn spawn( - channel: ChannelId, - handle: Handle, - boundary_exec: Arc, - login: bool, - ) -> Self { - let shared = Arc::new(Shared { - state: Mutex::new(State { - phase: Phase::AwaitBegin, - plan: None, - sender: None, - input_task: None, - progress: None, - process: None, - received: 0, - digest: Sha256::new(), - summary: None, - failure: None, - }), - changed: Notify::new(), - }); - let owner_shared = shared.clone(); - let owner = tokio::spawn(async move { - run(owner_shared, channel, handle, boundary_exec, login).await; - }); +impl Upload { + pub(super) fn new() -> Self { Self { - decoder: Decoder::default(), - shared, - owner, + plan: Some(PlanReader::default()), + received: 0, + progress: None, + process: None, + acks: None, + reported: Arc::new(AtomicBool::new(false)), } } - pub(super) fn data(&mut self, mut bytes: &[u8]) { - while !bytes.is_empty() { - match self.decoder.decode(&mut bytes) { - Ok(Some(frame)) => { - if let Err(failure) = self.accept(frame) { - self.fail(failure); - return; - } - } - Ok(None) => return, - Err(error) => { - self.fail(failure(FailureCode::Protocol, error.to_string())); - return; - } - } - } + pub(super) const fn awaiting_plan(&self) -> bool { + self.plan.is_some() } - pub(super) fn eof(&self) { - let valid = self.decoder.is_empty() - && self - .shared - .state - .lock() - .is_ok_and(|state| matches!(state.phase, Phase::Draining | Phase::Finished)); - if !valid { - self.fail(failure( - FailureCode::Protocol, - "upload stream ended before a complete END frame", - )); + /// Consume plan bytes, leaving archive bytes that follow it in `input`. + pub(super) fn read_plan(&mut self, input: &mut &[u8]) -> Result, PlanError> { + let plan = self + .plan + .as_mut() + .map_or(Ok(None), |reader| reader.read(input))?; + if plan.is_some() { + self.plan = None; } + Ok(plan) } - pub(super) fn reject_extended_data(&self) { - self.fail(failure( - FailureCode::Protocol, - "extended data is not supported by the upload subsystem", + /// Report extraction progress and return the exit task's finisher. + pub(super) fn start( + &mut self, + channel: ChannelId, + handle: Handle, + progress: watch::Receiver, + process: Arc, + ) -> Finisher { + let acks = tokio::spawn(send_acks( + handle, + channel, + progress.clone(), + process.clone(), )); + self.acks = Some(acks.abort_handle()); + self.progress = Some(progress.clone()); + self.process = Some(process); + Finisher { + channel, + progress, + acks, + reported: self.reported.clone(), + } } - pub(super) fn process(&self) -> Option> { - self.shared.state.lock().ok()?.process.clone() + /// Whether `bytes` more fit in the credit granted to the client. + pub(super) fn admits(&self, bytes: usize) -> bool { + let consumed = self + .progress + .as_ref() + .map_or(0, |progress| progress.borrow().consumed); + self.received + bytes as u64 - consumed <= WINDOW_BYTES } - fn fail(&self, failure: UploadFailure) { - if let Ok(mut state) = self.shared.state.lock() - && state.phase != Phase::Finished - && state.failure.is_none() - { - state.failure = Some(failure); - // Stop a stalled write immediately. The owner reports the typed - // failure and terminates the process outside the session callback. - state.sender.take(); - state.input_task.take(); - } - self.shared.changed.notify_one(); + pub(super) fn record(&mut self, bytes: usize) { + self.received += bytes as u64; } - fn accept(&self, frame: Frame) -> Result<(), UploadFailure> { - let mut state = self - .shared - .state - .lock() - .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; - if state.failure.is_some() || state.phase == Phase::Finished { - return Ok(()); - } - match frame { - Frame::Begin(plan) if state.phase == Phase::AwaitBegin => { - plan.validate() - .map_err(|error| failure(FailureCode::Protocol, error.to_string()))?; - state.phase = Phase::Starting; - state.plan = Some(plan); - } - Frame::Data(data) if state.phase == Phase::Receiving => { - let consumed = state - .progress - .as_ref() - .map_or(0, |progress| progress.borrow().consumed); - let received = state - .received - .checked_add(data.len() as u64) - .ok_or_else(|| failure(FailureCode::Protocol, "upload length overflow"))?; - if received - consumed > u64::from(WINDOW_BYTES) { - return Err(failure( - FailureCode::Protocol, - "upload input credit exceeded", - )); - } - let sender = state.sender.as_ref().ok_or_else(|| { - failure( - FailureCode::InputWrite, - "extraction process stdin is closed", - ) - })?; - sender.send(&data).map_err(|error| match error { - SendError::Full => { - failure(FailureCode::Internal, "upload input budget exceeded") - } - SendError::Closed => failure( - FailureCode::InputWrite, - "extraction process stdin is closed", - ), - })?; - state.received = received; - state.digest.update(data); - } - Frame::End(summary) if state.phase == Phase::Receiving => { - let actual = UploadSummary { - bytes: state.received, - sha256: state.digest.clone().finalize().into(), - }; - if summary != actual { - return Err(failure( - FailureCode::Protocol, - "upload count or digest mismatch", - )); - } - state.summary = Some(summary); - state.phase = Phase::Draining; - // EOF drains accepted bytes before closing boundary stdin. - state.sender.take(); - } - Frame::Cancel => { - return Err(failure(FailureCode::Protocol, "upload cancelled by client")); - } - _ => { - return Err(failure(FailureCode::Protocol, "unexpected upload frame")); - } - } - drop(state); - self.shared.changed.notify_one(); - Ok(()) + /// Claim the result for a rejection. False means the result was reported. + pub(super) fn reject(&self) -> bool { + !self.reported.swap(true, Ordering::AcqRel) } } -impl Drop for UploadChannel { +impl Drop for Upload { fn drop(&mut self) { - self.owner.abort(); - let process = self.shared.state.lock().ok().and_then(|mut state| { - state.phase = Phase::Finished; - state.sender.take(); - state.input_task.take(); - state.process.take() - }); - if let Some(process) = process { + if let Some(acks) = self.acks.take() { + acks.abort(); + } + // Disconnects drop channel state without a close request. + if self.reject() + && let Some(process) = self.process.take() + { tokio::spawn(async move { let _ = process.terminate().await; }); @@ -242,371 +122,70 @@ impl Drop for UploadChannel { } } -fn failure(code: FailureCode, message: impl Into) -> UploadFailure { - let mut message = message.into(); - // Error formatting and backend diagnostics can be large. Bound the UTF-8 - // string before it enters a queued protocol frame. - if message.len() > MAX_DIAGNOSTIC_BYTES { - let mut end = MAX_DIAGNOSTIC_BYTES; - while !message.is_char_boundary(end) { - end -= 1; - } - message.truncate(end); - } - UploadFailure { code, message } -} - -async fn send(handle: &Handle, channel: ChannelId, frame: Frame) -> bool { - match frame.encode() { - Ok(bytes) => handle.data(channel, bytes).await.is_ok(), - Err(_) => false, - } -} - -async fn finish( - shared: &Shared, - channel: ChannelId, - handle: &Handle, - result: Result, -) { - let (result, process) = shared.state.lock().map_or_else( - |_| { - ( - Err(failure(FailureCode::Internal, "upload state unavailable")), - None, - ) - }, - |mut state| { - // Commit the result under the same lock that records protocol failure. - let result = state.failure.take().map_or(result, Err); - state.phase = Phase::Finished; - state.sender.take(); - state.input_task.take(); - (result, state.process.take()) - }, - ); - let failed = result.is_err(); - if failed && let Some(process) = process { - tokio::spawn(async move { - let _ = process.terminate().await; - }); - } - let frame = match result { - Ok(summary) => Frame::Result(summary), - Err(error) => Frame::Error(error), - }; - let _ = send(handle, channel, frame).await; - let _ = handle - .exit_status_request(channel, if failed { 74 } else { 0 }) - .await; - let _ = handle.eof(channel).await; - let _ = handle.close(channel).await; -} - -async fn run( - shared: Arc, +pub(super) struct Finisher { channel: ChannelId, - handle: Handle, - boundary_exec: Arc, - login: bool, -) { - let result = drive(&shared, channel, &handle, boundary_exec, login).await; - finish(&shared, channel, &handle, result).await; -} - -async fn drive( - shared: &Shared, - channel: ChannelId, - handle: &Handle, - boundary_exec: Arc, - login: bool, -) -> Result { - if !send( - handle, - channel, - Frame::Hello { - window_bytes: WINDOW_BYTES, - max_data_bytes: u32::try_from(MAX_DATA_BYTES).expect("bounded data frame"), - }, - ) - .await - { - return Err(failure( - FailureCode::Internal, - "upload reply delivery failed", - )); - } - let plan = loop { - let changed = shared.changed.notified(); - { - let mut state = shared - .state - .lock() - .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; - if let Some(error) = state.failure.take() { - return Err(error); - } - if let Some(plan) = state.plan.take() { - break plan; - } - } - changed.await; - }; - let extraction = async { - let command = build_extract_command(&plan) - .map_err(|error| failure(FailureCode::Protocol, error.to_string()))?; - // Remote backends may admit a process before returning its attachment. - // Keep startup owned independently of this channel: cancellation must - // terminate even a result delivered after the upload owner is dropped. - let (started_tx, started_rx) = oneshot::channel(); - tokio::spawn(async move { - let result = boundary_exec - .exec(ExecSpec { - program: String::new(), - args: Vec::new(), - shell: Some(ShellSpec { - command: Some(command), - login, - }), - runtime_helper: None, - env: Vec::new(), - workdir: None, - pty: false, - }) - .await - .map(|exec| PendingExec(Some(exec))); - // A queued PendingExec also terminates on drop if the receiver was - // cancelled after send succeeded but before receiving the attachment. - let _ = started_tx.send(result); - }); - let mut pending = started_rx - .await - .map_err(|_| failure(FailureCode::Internal, "extraction startup stopped"))? - .map_err(|error| failure(FailureCode::Extract, error.to_string()))?; - { - let mut state = shared - .state - .lock() - .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; - // Channel drop can run on another worker after the receive completes. - // Keep the startup guard armed until ownership is registered under the - // same lock that cancellation uses to mark the channel finished. - if state.phase == Phase::Finished || state.failure.is_some() { - return Err(failure( - FailureCode::Protocol, - "upload cancelled during startup", - )); - } - state.process = Some( - pending - .0 - .as_ref() - .expect("owned startup result") - .process - .clone(), + progress: watch::Receiver, + acks: JoinHandle<()>, + reported: Arc, +} + +impl Finisher { + /// Send the final total for a successful extraction and return the exit + /// status to report, or `None` when a rejection already reported it. + pub(super) async fn finish(mut self, handle: &Handle, code: u32) -> Option { + self.acks.abort(); + let _ = (&mut self.acks).await; + let drained = code == 0 + && matches!( + tokio::time::timeout( + DRAIN_TIMEOUT, + self.progress.wait_for(|progress| progress.state != InputState::Writing), + ) + .await, + Ok(Ok(progress)) if progress.state == InputState::Drained ); + if self.reported.swap(true, Ordering::AcqRel) { + return None; } - let mut exec = pending.0.take().expect("registered startup result"); - let stdin = exec - .stdin - .take() - .ok_or_else(|| failure(FailureCode::Internal, "extraction process has no stdin"))?; - let (sender, input_task, mut progress) = InputTask::spawn_with_progress(stdin); - { - let mut state = shared - .state - .lock() - .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; - if state.phase == Phase::Finished { - return Err(failure( - FailureCode::Protocol, - "upload cancelled during startup", - )); - } - if let Some(error) = state.failure.take() { - return Err(error); - } - state.sender = Some(sender); - state.input_task = Some(input_task); - state.progress = Some(progress.clone()); - state.phase = Phase::Receiving; - } - let diagnostic = Arc::new(Mutex::new(Vec::new())); - let stdout = DrainTask::spawn(exec.stdout, diagnostic.clone()); - let stderr = exec - .stderr - .map(|stderr| DrainTask::spawn(stderr, diagnostic.clone())); - let mut exit = Box::pin(exec.process.wait()); - if !send(handle, channel, Frame::Ready).await { - return Err(failure( - FailureCode::Internal, - "upload reply delivery failed", - )); - } - let mut acknowledged = 0; - let status = loop { - let latest = *progress.borrow_and_update(); - if latest.state == InputState::Failed { - return Err(failure( - FailureCode::InputWrite, - diagnostic_message(&diagnostic, "failed to write extraction stdin"), - )); - } - if latest.consumed - acknowledged >= ACK_INTERVAL - || (latest.state == InputState::Drained && latest.consumed != acknowledged) - { - if !send(handle, channel, Frame::Ack(latest.consumed)).await { - return Err(failure( - FailureCode::Internal, - "upload reply delivery failed", - )); - } - acknowledged = latest.consumed; - } - tokio::select! { - status = &mut exit => { - let status = status.map_err(|error| failure(FailureCode::Extract, error.to_string()))?; - let final_progress = *progress.borrow(); - let (final_summary, final_received) = shared.state.lock().map(|state| (state.summary, state.received)).map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))?; - if final_summary.is_none() || final_progress.consumed != final_received { - return Err(extract_failure(status, &diagnostic_message(&diagnostic, "extraction process exited before upload completed"))); - } - // A pipe shutdown can wake the child just before the writer - // publishes Drained. Await that final watch update rather than - // treating the successful EOF race as an early extraction exit. - while progress.borrow().state == InputState::Writing { - progress.changed().await.map_err(|_| failure(FailureCode::InputWrite, "extraction stdin writer stopped"))?; - } - if progress.borrow().state != InputState::Drained { - return Err(failure(FailureCode::InputWrite, diagnostic_message(&diagnostic, "failed to close extraction stdin"))); - } - if final_received != acknowledged && !send(handle, channel, Frame::Ack(final_received)).await { - return Err(failure(FailureCode::Internal, "upload reply delivery failed")); - } - break status; - } - result = progress.changed() => { - if result.is_err() { - return Err(failure(FailureCode::InputWrite, "extraction stdin writer stopped")); - } - } - } - }; - stdout - .join() - .await - .map_err(|error| failure(FailureCode::Extract, error.to_string()))?; - if let Some(stderr) = stderr { - stderr - .join() - .await - .map_err(|error| failure(FailureCode::Extract, error.to_string()))?; - } - let status = if let Some(output_status) = exec.output_status { - output_status - .await - .map_err(|_| failure(FailureCode::Extract, "extraction output delivery failed"))? - } else { - status - }; - if status != BoundaryExitStatus::Exited(0) { - let diagnostic = diagnostic_message(&diagnostic, "extraction process failed"); - return Err(extract_failure(status, &diagnostic)); + if code != 0 { + return Some(code); } - shared - .state - .lock() - .ok() - .and_then(|state| state.summary) - .ok_or_else(|| failure(FailureCode::Internal, "upload summary unavailable")) - }; - tokio::pin!(extraction); - loop { - let changed = shared.changed.notified(); - let error = shared - .state - .lock() - .map_err(|_| failure(FailureCode::Internal, "upload state unavailable"))? - .failure - .clone(); - if let Some(error) = error { - return Err(error); - } - tokio::select! { - result = &mut extraction => return result, - () = changed => {} + if !drained { + return Some(74); } + let consumed = self.progress.borrow().consumed; + let _ = handle.data(self.channel, ack(consumed)).await; + Some(0) } } -fn diagnostic_message(diagnostic: &Mutex>, fallback: &str) -> String { - diagnostic - .lock() - .ok() - .filter(|bytes| !bytes.is_empty()) - .map_or_else( - || fallback.to_string(), - |bytes| String::from_utf8_lossy(&bytes).into_owned(), - ) -} - -fn extract_failure(status: BoundaryExitStatus, diagnostic: &str) -> UploadFailure { - if status == BoundaryExitStatus::Exited(DESTINATION_CHANGED_EXIT_CODE) { - failure(FailureCode::DestinationChanged, diagnostic) - } else { - failure( - FailureCode::Extract, - format!("extraction failed ({status:?}): {diagnostic}"), - ) - } +fn ack(consumed: u64) -> Vec { + consumed.to_be_bytes().to_vec() } -/// Own a late startup result until the channel registers process ownership. -struct PendingExec(Option); - -impl Drop for PendingExec { - fn drop(&mut self) { - if let Some(exec) = self.0.take() { - tokio::spawn(async move { - let _ = exec.process.terminate().await; - }); +async fn send_acks( + handle: Handle, + channel: ChannelId, + mut progress: watch::Receiver, + process: Arc, +) { + // Zero tells the client that extraction started. + if handle.data(channel, ack(0)).await.is_err() { + return; + } + let mut acknowledged = 0; + while progress.changed().await.is_ok() { + let latest = *progress.borrow_and_update(); + if latest.state == InputState::Failed { + // The extractor stopped reading; its exit status explains why. + let _ = process.terminate().await; + return; } - } -} - -/// Aborting the owner also aborts output drains; dropping `JoinHandle` alone -/// would detach a pipe reader and retain the boundary streams after close. -struct DrainTask(JoinHandle>); - -impl DrainTask { - fn spawn( - mut stream: Box, - diagnostic: Arc>>, - ) -> Self { - Self(tokio::spawn(async move { - let mut bytes = [0; 4096]; - loop { - match stream.read(&mut bytes).await { - Ok(0) => return Ok(()), - Err(error) => return Err(error), - Ok(size) => { - if let Ok(mut diagnostic) = diagnostic.lock() { - let retain = size.min(MAX_DIAGNOSTIC_BYTES - diagnostic.len()); - diagnostic.extend_from_slice(&bytes[..retain]); - } - } - } + if latest.consumed - acknowledged >= ACK_INTERVAL { + if handle.data(channel, ack(latest.consumed)).await.is_err() { + return; } - })) - } - - async fn join(mut self) -> std::io::Result<()> { - (&mut self.0).await.map_err(std::io::Error::other)? - } -} - -impl Drop for DrainTask { - fn drop(&mut self) { - self.0.abort(); + acknowledged = latest.consumed; + } } } diff --git a/crates/openshell-supervisor-process/src/ssh/upload_tests.rs b/crates/openshell-supervisor-process/src/ssh/upload_tests.rs index f9cd4c4ef5..b49cb6bf77 100644 --- a/crates/openshell-supervisor-process/src/ssh/upload_tests.rs +++ b/crates/openshell-supervisor-process/src/ssh/upload_tests.rs @@ -3,16 +3,13 @@ //! Exercise the production upload handler over a real SSH session. -use super::tests::{AcceptAnyServerKey, test_client}; -use openshell_core::upload::{ - Decoder, FailureCode, Frame, MAX_DATA_BYTES, SUBSYSTEM, UploadPlan, UploadSummary, WINDOW_BYTES, -}; +use super::tests::{AcceptAnyServerKey, RejectingExec, test_client}; +use openshell_core::upload::{SUBSYSTEM, UploadPlan, WINDOW_BYTES}; use openshell_isolation_interface::contract::{ BackendError, BoundaryExec, BoundaryExitStatus, BoundaryProcess, BoundarySignal, ExecSession, ExecSpec, }; use russh::ChannelMsg; -use sha2::{Digest, Sha256}; use std::collections::VecDeque; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -20,6 +17,8 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt, DuplexStream}; use tokio::sync::{mpsc, oneshot, watch}; const DEADLINE: Duration = Duration::from_secs(10); +const CHUNK: usize = 64 * 1024; +const WINDOW: usize = 1024 * 1024; struct Process { status: watch::Sender>, @@ -78,6 +77,16 @@ struct Control { signals: mpsc::UnboundedReceiver, } +impl Process { + async fn terminated(&self) { + self.terminated + .subscribe() + .wait_for(|terminated| *terminated) + .await + .unwrap(); + } +} + fn exec_session() -> (ExecSession, Control) { let (stdin, reader) = tokio::io::duplex(1024); let (stdout, writer) = tokio::io::duplex(4096); @@ -108,100 +117,100 @@ fn exec_session() -> (ExecSession, Control) { ) } +async fn client_for(sessions: Vec) -> russh::client::Handle { + test_client( + None, + Arc::new(Exec(Mutex::new(sessions.into()))), + russh::client::Config::default(), + ) + .await +} + type Channel = russh::Channel; +/// Collects acknowledgment totals from stdout, stderr and exit statuses. +#[derive(Default)] struct Reader { - decoder: Decoder, - frames: VecDeque, + stdout: Vec, + stderr: Vec, statuses: Vec, } impl Reader { - fn new() -> Self { - Self { - decoder: Decoder::default(), - frames: VecDeque::new(), - statuses: Vec::new(), - } - } - - async fn frame(&mut self, channel: &mut Channel) -> Frame { + /// Next acknowledged total, or `None` once the channel closes. + async fn next(&mut self, channel: &mut Channel) -> Option { loop { - if let Some(frame) = self.frames.pop_front() { - return frame; + if self.stdout.len() >= 8 { + let total = self.stdout.drain(..8).collect::>(); + return Some(u64::from_be_bytes(total.try_into().unwrap())); } - match channel.wait().await.expect("channel closed before frame") { - ChannelMsg::Data { data } => { - let mut bytes = data.as_ref(); - while !bytes.is_empty() { - if let Some(frame) = self.decoder.decode(&mut bytes).unwrap() { - self.frames.push_back(frame); - } - } - } + match channel.wait().await? { + ChannelMsg::Data { data } => self.stdout.extend_from_slice(&data), + ChannelMsg::ExtendedData { data, ext: 1 } => self.stderr.extend_from_slice(&data), ChannelMsg::ExitStatus { exit_status } => self.statuses.push(exit_status), + ChannelMsg::Close => return None, ChannelMsg::WindowAdjusted { .. } | ChannelMsg::Success | ChannelMsg::Eof => {} message => panic!("unexpected SSH message: {message:?}"), } } } + + /// Remaining totals through channel close; no partial total may remain. + async fn close(&mut self, channel: &mut Channel) -> Vec { + let mut totals = Vec::new(); + while let Some(total) = self.next(channel).await { + totals.push(total); + } + assert!(self.stdout.is_empty(), "partial acknowledgment"); + totals + } } -async fn send(channel: &Channel, frame: Frame) { - channel - .data(frame.encode().unwrap().as_slice()) - .await - .unwrap(); +fn plan_line() -> Vec { + UploadPlan { + destination_directory: "/tmp/upload".to_string(), + expected_destination: None, + } + .encode() + .unwrap() } -async fn start(client: &russh::client::Handle) -> (Channel, Reader) { +async fn open(client: &russh::client::Handle) -> Channel { let mut channel = client.channel_open_session().await.unwrap(); channel.request_subsystem(true, SUBSYSTEM).await.unwrap(); assert!(matches!(channel.wait().await, Some(ChannelMsg::Success))); - let mut reader = Reader::new(); - assert_eq!( - reader.frame(&mut channel).await, - Frame::Hello { - window_bytes: WINDOW_BYTES, - max_data_bytes: u32::try_from(MAX_DATA_BYTES).unwrap(), - } - ); - send( - &channel, - Frame::Begin(UploadPlan { - destination_directory: "/tmp/upload".to_string(), - expected_destination: None, - }), - ) - .await; - assert_eq!(reader.frame(&mut channel).await, Frame::Ready); + channel +} + +async fn start(client: &russh::client::Handle) -> (Channel, Reader) { + let mut channel = open(client).await; + channel.data(plan_line().as_slice()).await.unwrap(); + let mut reader = Reader::default(); + assert_eq!(reader.next(&mut channel).await, Some(0)); (channel, reader) } #[tokio::test] async fn stalled_upload_preserves_controls_then_resumes_beyond_four_mib() { + assert_eq!(WINDOW as u64, WINDOW_BYTES); tokio::time::timeout(DEADLINE, async { let (upload, mut control) = exec_session(); let (sibling, mut other) = exec_session(); - let client = test_client( - None, - Arc::new(Exec(Mutex::new([upload, sibling].into()))), - russh::client::Config::default(), - ) - .await; + let client = client_for(vec![upload, sibling]).await; let (mut channel, mut reader) = start(&client).await; - let payload = vec![42; MAX_DATA_BYTES]; - let mut sent = 0_u64; - while sent < u64::from(WINDOW_BYTES) { - send(&channel, Frame::Data(payload.clone())).await; - sent += payload.len() as u64; + let chunk = vec![42; CHUNK]; + let mut sent = 0; + while sent < WINDOW { + channel.data(chunk.as_slice()).await.unwrap(); + sent += CHUNK; } channel.signal(russh::Sig::INT).await.unwrap(); assert_eq!(control.signals.recv().await, Some(BoundarySignal::Int)); assert!( - tokio::time::timeout(Duration::from_millis(50), reader.frame(&mut channel)) + tokio::time::timeout(Duration::from_millis(50), reader.next(&mut channel)) .await - .is_err() + .is_err(), + "credit returned before the extractor consumed input" ); // Stalled upload input must leave sibling channel requests responsive. let mut sibling = client.channel_open_session().await.unwrap(); @@ -213,9 +222,8 @@ async fn stalled_upload_preserves_controls_then_resumes_beyond_four_mib() { other.stdin.read_to_end(&mut sibling_input).await.unwrap(); assert_eq!(sibling_input, b"sibling"); sibling.close().await.unwrap(); - // Output is drained into a bounded diagnostic, never exposed as forged - // protocol stdout or allowed to block extraction on a full pipe. - let forged = Frame::Ack(u64::MAX).encode().unwrap(); + // Extractor stdout is delivered as stderr, so it cannot forge credit. + let forged = u64::MAX.to_be_bytes(); control.stdout.write_all(&forged).await.unwrap(); let noisy = tokio::spawn(async move { control @@ -223,123 +231,118 @@ async fn stalled_upload_preserves_controls_then_resumes_beyond_four_mib() { .write_all(&vec![9; 1024 * 1024]) .await .unwrap(); - drop(control.stderr); }); let consumed = tokio::spawn(async move { let mut bytes = Vec::new(); control.stdin.read_to_end(&mut bytes).await.unwrap(); - assert_eq!(bytes, vec![42; 5 * 1024 * 1024]); + assert_eq!(bytes, vec![42; 5 * WINDOW]); drop(control.stdout); control .process .status .send_replace(Some(BoundaryExitStatus::Exited(0))); }); - let mut ack = 0; - while sent < 5 * 1024 * 1024 { - if sent + MAX_DATA_BYTES as u64 > ack + u64::from(WINDOW_BYTES) { - match reader.frame(&mut channel).await { - Frame::Ack(consumed) => { - assert!(consumed >= ack && consumed <= sent); - ack = consumed; - } - frame => panic!("unexpected transfer reply: {frame:?}"), - } + let mut acknowledged = 0; + while sent < 5 * WINDOW { + if sent + CHUNK > acknowledged + WINDOW { + let total = reader.next(&mut channel).await.unwrap(); + assert!(total >= acknowledged as u64 && total <= sent as u64); + acknowledged = usize::try_from(total).unwrap(); } else { - send(&channel, Frame::Data(payload.clone())).await; - sent += payload.len() as u64; + channel.data(chunk.as_slice()).await.unwrap(); + sent += CHUNK; } } - let summary = UploadSummary { - bytes: sent, - sha256: Sha256::digest(vec![42; usize::try_from(sent).unwrap()]).into(), - }; - send(&channel, Frame::End(summary)).await; channel.eof().await.unwrap(); - loop { - match reader.frame(&mut channel).await { - Frame::Ack(consumed) => assert!(consumed <= sent), - Frame::Result(actual) => { - assert_eq!(actual, summary); - break; - } - frame => panic!("unexpected terminal reply: {frame:?}"), - } - } - let mut statuses = Vec::new(); - while let Some(message) = channel.wait().await { - match message { - ChannelMsg::ExitStatus { exit_status } => statuses.push(exit_status), - ChannelMsg::Close => break, - _ => {} - } - } - assert_eq!(statuses, [0]); - consumed.await.unwrap(); noisy.await.unwrap(); + let totals = reader.close(&mut channel).await; + assert_eq!(totals.last(), Some(&(sent as u64))); + assert_eq!(reader.statuses, [0]); + assert_eq!(reader.stderr.len(), forged.len() + 1024 * 1024); + assert!(reader.stderr.windows(8).any(|bytes| bytes == forged)); + consumed.await.unwrap(); }) .await .expect("upload did not resume or blocked SSH control traffic"); } #[tokio::test] -async fn credit_violation_and_truncated_eof_are_typed_and_terminate() { - for truncated in [false, true] { +async fn invalid_plans_and_credit_violations_exit_74() { + for case in ["credit", "eof", "plan"] { tokio::time::timeout(DEADLINE, async { let (session, control) = exec_session(); - let client = test_client( - None, - Arc::new(Exec(Mutex::new([session].into()))), - russh::client::Config::default(), - ) - .await; - let (mut channel, mut reader) = start(&client).await; - if truncated { - channel.eof().await.unwrap(); + let client = client_for(vec![session]).await; + let (mut channel, mut reader) = if case == "credit" { + start(&client).await } else { - for _ in 0..=WINDOW_BYTES as usize / MAX_DATA_BYTES { - send(&channel, Frame::Data(vec![7; MAX_DATA_BYTES])).await; + (open(&client).await, Reader::default()) + }; + match case { + "credit" => { + for _ in 0..=WINDOW / CHUNK { + channel.data(vec![7; CHUNK].as_slice()).await.unwrap(); + } } + "eof" => { + channel.data(&b"{"[..]).await.unwrap(); + channel.eof().await.unwrap(); + } + _ => { + let plan = br#"{"destination_directory":"/tmp","expected_destination":null,"command":"id"} +"#; + channel.data(&plan[..]).await.unwrap(); + } + } + reader.close(&mut channel).await; + assert_eq!(reader.statuses, [74]); + if case == "credit" { + control.process.terminated().await; + } else { + let stderr = String::from_utf8_lossy(&reader.stderr); + assert!(stderr.contains("openshell: "), "{stderr}"); + assert!(!*control.process.terminated.borrow()); } - assert!(matches!(reader.frame(&mut channel).await, Frame::Error(error) if error.code == FailureCode::Protocol)); - control.process.terminated.subscribe().wait_for(|terminated| *terminated).await.unwrap(); - let mut eof = Vec::new(); - let _ = control.stdin.take(2).read_to_end(&mut eof).await; }) .await - .expect("protocol failure retained a stalled input or process"); + .expect("rejected upload did not close"); } } #[tokio::test] -async fn broken_stdin_and_early_exit_wake_credit_waiter() { - for early_exit in [false, true] { +async fn extractor_failures_report_status_and_output() { + for early_exit in [true, false] { tokio::time::timeout(DEADLINE, async { let (session, mut control) = exec_session(); - let client = test_client( - None, - Arc::new(Exec(Mutex::new([session].into()))), - russh::client::Config::default(), - ) - .await; + let client = client_for(vec![session]).await; let (mut channel, mut reader) = start(&client).await; let diagnostic = "mkdir: /tmp/existing-file/child: Not a directory"; - let mut stderr = diagnostic.as_bytes().to_vec(); - stderr.resize(8192, b' '); - // Filling more than the pipe holds proves the drain captured the cause. - control.stderr.write_all(&stderr).await.unwrap(); if early_exit { - control.process.status.send_replace(Some(BoundaryExitStatus::Exited(1))); + control + .stderr + .write_all(diagnostic.as_bytes()) + .await + .unwrap(); + control + .process + .status + .send_replace(Some(BoundaryExitStatus::Exited(2))); } else { + // Broken stdin wakes no credit; the extractor is stopped. drop(control.stdin); - send(&channel, Frame::Data(vec![7; MAX_DATA_BYTES])).await; + channel.data(vec![7; CHUNK].as_slice()).await.unwrap(); + control.process.terminated().await; + } + drop((control.stdout, control.stderr)); + reader.close(&mut channel).await; + if early_exit { + assert_eq!(reader.statuses, [2]); + assert_eq!(reader.stderr, diagnostic.as_bytes()); + } else { + assert_eq!(reader.statuses, [74]); } - let expected = if early_exit { FailureCode::Extract } else { FailureCode::InputWrite }; - assert!(matches!(reader.frame(&mut channel).await, Frame::Error(error) if error.code == expected && error.message.contains(diagnostic))); - control.process.terminated.subscribe().wait_for(|terminated| *terminated).await.unwrap(); }) .await - .expect("terminal failure did not wake upload sender"); + .expect("extractor failure did not close the upload"); } } @@ -348,14 +351,9 @@ async fn close_and_disconnect_cancel_stalled_boundary_input() { for disconnect in [false, true] { tokio::time::timeout(DEADLINE, async { let (session, mut control) = exec_session(); - let client = test_client( - None, - Arc::new(Exec(Mutex::new([session].into()))), - russh::client::Config::default(), - ) - .await; + let client = client_for(vec![session]).await; let (channel, _) = start(&client).await; - send(&channel, Frame::Data(vec![7; MAX_DATA_BYTES])).await; + channel.data(vec![7; CHUNK].as_slice()).await.unwrap(); channel.signal(russh::Sig::INT).await.unwrap(); control.signals.recv().await.unwrap(); if disconnect { @@ -366,13 +364,7 @@ async fn close_and_disconnect_cancel_stalled_boundary_input() { } else { channel.close().await.unwrap(); } - control - .process - .terminated - .subscribe() - .wait_for(|terminated| *terminated) - .await - .unwrap(); + control.process.terminated().await; let mut received = Vec::new(); control.stdin.read_to_end(&mut received).await.unwrap(); assert!(received.len() <= 1024); @@ -387,7 +379,7 @@ async fn upload_rejects_pty_and_duplicate_channel_requests() { tokio::time::timeout(DEADLINE, async { let client = test_client( None, - Arc::new(super::tests::RejectingExec), + Arc::new(RejectingExec), russh::client::Config::default(), ) .await; @@ -398,14 +390,7 @@ async fn upload_rejects_pty_and_duplicate_channel_requests() { assert!(matches!(pty.wait().await, Some(ChannelMsg::Success))); pty.request_subsystem(true, SUBSYSTEM).await.unwrap(); assert!(matches!(pty.wait().await, Some(ChannelMsg::Failure))); - let mut upload = client.channel_open_session().await.unwrap(); - upload.request_subsystem(true, SUBSYSTEM).await.unwrap(); - assert!(matches!(upload.wait().await, Some(ChannelMsg::Success))); - let mut reader = Reader::new(); - assert!(matches!( - reader.frame(&mut upload).await, - Frame::Hello { .. } - )); + let mut upload = open(&client).await; upload.exec(true, "cat").await.unwrap(); assert!(matches!(upload.wait().await, Some(ChannelMsg::Failure))); upload.request_shell(true).await.unwrap(); @@ -456,27 +441,15 @@ async fn cancellation_during_remote_startup_terminates_late_process() { russh::client::Config::default(), ) .await; - let mut channel = client.channel_open_session().await.unwrap(); - channel.request_subsystem(true, SUBSYSTEM).await.unwrap(); - assert!(matches!(channel.wait().await, Some(ChannelMsg::Success))); - let mut reader = Reader::new(); - assert!(matches!( - reader.frame(&mut channel).await, - Frame::Hello { .. } - )); - send( - &channel, - Frame::Begin(UploadPlan { - destination_directory: "/tmp/upload".into(), - expected_destination: None, - }), - ) - .await; + let channel = open(&client).await; + channel.data(plan_line().as_slice()).await.unwrap(); started .subscribe() .wait_for(|started| *started) .await .unwrap(); + // Startup holds the session loop, as for exec requests, so the + // close or disconnect is handled after the process is registered. if disconnect { client .disconnect(russh::Disconnect::ByApplication, "test", "") @@ -485,27 +458,8 @@ async fn cancellation_during_remote_startup_terminates_late_process() { } else { channel.close().await.unwrap(); } - // A signal on another channel is a session-loop barrier proving - // that the close was handled while startup is still pending. - // Disconnect has no surviving channel, so wait for client closure. - if disconnect { - let _ = client.await; - } else { - let mut sibling = client.channel_open_session().await.unwrap(); - sibling - .request_subsystem(true, "unsupported") - .await - .unwrap(); - assert!(matches!(sibling.wait().await, Some(ChannelMsg::Failure))); - } release.send(()).unwrap(); - control - .process - .terminated - .subscribe() - .wait_for(|terminated| *terminated) - .await - .unwrap(); + control.process.terminated().await; }) .await .expect("cancelled startup stranded a late extractor"); @@ -513,65 +467,56 @@ async fn cancellation_during_remote_startup_terminates_late_process() { } #[tokio::test] -async fn mismatched_end_and_extended_data_fail_without_result() { - for (extended, after_exit, pending_output_status) in [ - (false, false, false), - (true, false, false), - (true, true, false), - (true, true, true), - ] { +async fn late_rejection_wins_until_the_result_is_reported() { + for (after_exit, pending_output_status) in [(false, false), (true, false), (true, true)] { tokio::time::timeout(DEADLINE, async { let (mut session, mut control) = exec_session(); let (output_status, receiver) = oneshot::channel(); if pending_output_status { session.output_status = Some(receiver); } - let client = test_client(None, Arc::new(Exec(Mutex::new([session].into()))), russh::client::Config::default()).await; + let client = client_for(vec![session]).await; let (mut channel, mut reader) = start(&client).await; if after_exit { - let summary = UploadSummary { bytes: 1, sha256: Sha256::digest(b"x").into() }; - send(&channel, Frame::Data(b"x".to_vec())).await; - send(&channel, Frame::End(summary)).await; + channel.data(&b"x"[..]).await.unwrap(); + channel.eof().await.unwrap(); let mut bytes = Vec::new(); control.stdin.read_to_end(&mut bytes).await.unwrap(); assert_eq!(bytes, b"x"); - assert_eq!(reader.frame(&mut channel).await, Frame::Ack(1)); - control.process.status.send_replace(Some(BoundaryExitStatus::Exited(0))); - control.process.waited.subscribe().wait_for(|waited| *waited).await.unwrap(); if pending_output_status { + // Output EOF leaves the exit task waiting for the backend's + // final status; dropped readers prove both drains finished. control.stdout.shutdown().await.unwrap(); control.stderr.shutdown().await.unwrap(); assert_eq!(control.stdout.read(&mut [0]).await.unwrap(), 0); assert_eq!(control.stderr.read(&mut [0]).await.unwrap(), 0); - tokio::task::yield_now().await; - } - } - if extended { - channel.extended_data(1, b"unexpected".as_slice()).await.unwrap(); - } else { - send(&channel, Frame::End(UploadSummary {bytes: 1, sha256: [0;32]})).await; - } - assert!(matches!(reader.frame(&mut channel).await, Frame::Error(failure) if failure.code == FailureCode::Protocol)); - control.process.terminated.subscribe().wait_for(|terminated| *terminated).await.unwrap(); - // Error and cleanup must arrive while output EOF/status is still held. - if after_exit && !pending_output_status { - assert_eq!(control.stdout.read(&mut [0]).await.unwrap(), 0); - assert_eq!(control.stderr.read(&mut [0]).await.unwrap(), 0); - } - if pending_output_status { - assert!(output_status.is_closed()); - } - assert!(reader.frames.is_empty()); - while let Some(message) = channel.wait().await { - match message { - ChannelMsg::ExitStatus { exit_status } => reader.statuses.push(exit_status), - ChannelMsg::Data { .. } => panic!("unexpected upload frame after ERROR"), - ChannelMsg::Close => break, - _ => {} + } else { + control + .process + .status + .send_replace(Some(BoundaryExitStatus::Exited(0))); + control + .process + .waited + .subscribe() + .wait_for(|waited| *waited) + .await + .unwrap(); } } + // The exit task is still waiting for output EOF or its final + // status, so the rejection must report without waiting for it. + channel + .extended_data(1, b"unexpected".as_slice()) + .await + .unwrap(); + assert!(reader.close(&mut channel).await.is_empty()); assert_eq!(reader.statuses, [74]); - }).await.expect("invalid upload input did not terminate"); + control.process.terminated().await; + drop((control.stdout, control.stderr, output_status)); + }) + .await + .expect("invalid upload input did not terminate"); } } @@ -606,20 +551,12 @@ async fn upload_preserves_login_shell_selection() { } channel.request_subsystem(true, SUBSYSTEM).await.unwrap(); assert!(matches!(channel.wait().await, Some(ChannelMsg::Success))); - let mut reader = Reader::new(); - assert!(matches!( - reader.frame(&mut channel).await, - Frame::Hello { .. } - )); - send( - &channel, - Frame::Begin(UploadPlan { - destination_directory: "/tmp/upload".into(), - expected_destination: None, - }), - ) - .await; - assert!(matches!(reader.frame(&mut channel).await, Frame::Error(_))); + channel.data(plan_line().as_slice()).await.unwrap(); + let mut reader = Reader::default(); + assert!(reader.close(&mut channel).await.is_empty()); + assert_eq!(reader.statuses, [74]); + let stderr = String::from_utf8_lossy(&reader.stderr); + assert!(stderr.contains("login selection observed"), "{stderr}"); assert_eq!(*login.borrow(), Some(!no_login)); }) .await From be78f0883249708d72459522c117bd1899f8beea Mon Sep 17 00:00:00 2001 From: Shiju Date: Tue, 13 Oct 2026 12:47:10 +0530 Subject: [PATCH 4/5] fix(cli): isolate upload SSH connections and pin credit cadence - Apply the forward multiplexing overrides to the upload ssh command, so a user ControlMaster cannot carry an upload over another session's master connection. - Kill ssh before closing its stdin when a transfer fails locally, and include ssh stderr in transport errors. - Share ACK_INTERVAL_BYTES from core. Assert at compile time that CLI chunks fit the credit left after an ack and that the window stays below the supervisor's pending-input guard. - Log the real reason when the supervisor rejects an upload. - Test that a single-packet upload reports its final total. Signed-off-by: Shiju --- crates/openshell-cli/src/ssh.rs | 106 ++++++++++-------- crates/openshell-cli/src/ssh/upload.rs | 69 +++++++----- crates/openshell-core/src/upload.rs | 6 + crates/openshell-supervisor-process/README.md | 4 +- .../openshell-supervisor-process/src/ssh.rs | 4 +- .../src/ssh/upload.rs | 7 +- .../src/ssh/upload_tests.rs | 32 +++++- 7 files changed, 144 insertions(+), 84 deletions(-) diff --git a/crates/openshell-cli/src/ssh.rs b/crates/openshell-cli/src/ssh.rs index 1e50ef574e..02fe6515ed 100644 --- a/crates/openshell-cli/src/ssh.rs +++ b/crates/openshell-cli/src/ssh.rs @@ -69,9 +69,6 @@ const CONNECT_CHILD_TERMINATION_TIMEOUT: Duration = Duration::from_secs(5); const SSH_TRANSPORT_FAILURE_EXIT_CODE: i32 = 255; const SYNC_RETRY_ATTEMPTS: usize = 4; const SYNC_RETRY_DELAY: Duration = Duration::from_secs(2); -#[cfg(test)] -const UPLOAD_DESTINATION_CHANGED_EXIT_CODE: i32 = - openshell_core::upload::DESTINATION_CHANGED_EXIT_CODE; static DEFAULT_SSH_LOG_LEVEL: OnceLock<&'static str> = OnceLock::new(); /// Set the SSH log level derived from CLI verbosity before starting commands. @@ -825,10 +822,9 @@ pub async fn sandbox_forward( Ok(()) } -/// A forward must stay owned by the process we spawn so it can be reaped or -/// tracked by PID. User SSH config may otherwise move it into a persistent mux. -fn ssh_forward_command(proxy_command: &str, sandbox_id: &str, spec: &ForwardSpec) -> Command { - let mut command = ssh_base_command(proxy_command); +/// Keep the SSH connection owned by the process we spawn. User SSH config may +/// otherwise move it into a persistent mux. +fn own_ssh_connection(command: &mut Command) -> &mut Command { command .arg("-o") .arg("ControlMaster=no") @@ -838,6 +834,13 @@ fn ssh_forward_command(proxy_command: &str, sandbox_id: &str, spec: &ForwardSpec .arg("ControlPersist=no") .arg("-o") .arg("ForkAfterAuthentication=no") +} + +/// A forward must stay owned by the process we spawn so it can be reaped or +/// tracked by PID. +fn ssh_forward_command(proxy_command: &str, sandbox_id: &str, spec: &ForwardSpec) -> Command { + let mut command = ssh_base_command(proxy_command); + own_ssh_connection(&mut command) .arg("-N") .arg("-o") .arg("ExitOnForwardFailure=yes") @@ -1307,13 +1310,7 @@ async fn ssh_tar_upload( UploadSource::FileList { .. } => dest_dir.to_string(), }; - let mut ssh = ssh_base_command(&session.proxy_command); - ssh.arg("-T") - .arg("-o") - .arg("RequestTTY=no") - .arg("-s") - .arg("sandbox") - .arg(openshell_core::upload::SUBSYSTEM); + let ssh = ssh_upload_command(&session.proxy_command); let plan = UploadPlan { destination_directory: dest_dir.to_string(), expected_destination: expected_destination.map(|(path, kind)| ExpectedDestination { @@ -1330,6 +1327,20 @@ async fn ssh_tar_upload( Ok(final_path) } +/// Upload EOF marks a complete archive. Under a user SSH mux, killing the client +/// would still let the master send EOF, so the upload owns its connection. +fn ssh_upload_command(proxy_command: &str) -> Command { + let mut command = ssh_base_command(proxy_command); + own_ssh_connection(&mut command) + .arg("-T") + .arg("-o") + .arg("RequestTTY=no") + .arg("-s") + .arg("sandbox") + .arg(openshell_core::upload::SUBSYSTEM); + command +} + /// Split a sandbox path into (`parent_directory`, basename). /// /// Examples: @@ -2506,7 +2517,7 @@ mod tests { fn upload_errors_retry_only_explicit_transport_failures() { for error in [ upload::UploadError::Unsupported, - upload::UploadError::Protocol("connection reset in malformed frame".into()), + upload::UploadError::Protocol("connection reset in invalid acknowledgment".into()), upload::UploadError::Remote("destination /connection/unavailable".into()), upload::UploadError::Archive("broken pipe while reading source".into()), upload::UploadError::Cancelled, @@ -2852,7 +2863,7 @@ mod tests { #[cfg(unix)] #[test] - fn forward_ssh_command_overrides_user_multiplexing_and_forking() { + fn forward_and_upload_ssh_commands_override_user_multiplexing_and_forking() { let _guard = TEST_ENV_LOCK .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); @@ -2867,36 +2878,37 @@ mod tests { ) .unwrap(); - let forward = ssh_forward_command( - "openshell ssh-proxy --sandbox demo", - "sandbox-id", - &ForwardSpec::new(18080), - ); - let output = Command::new("ssh") - .arg("-F") - .arg(&config) - .arg("-G") - .args(forward.get_args()) - .output() - .unwrap(); - assert!( - output.status.success(), - "ssh -G failed: {}", - String::from_utf8_lossy(&output.stderr) - ); - let effective = String::from_utf8(output.stdout).unwrap(); - assert!(effective.lines().any(|line| line == "controlmaster false")); - assert!(effective.lines().any(|line| line == "controlpersist no")); - assert!( - effective - .lines() - .any(|line| line == "forkafterauthentication no") - ); - assert!( - !effective - .lines() - .any(|line| { line.starts_with("controlpath ") && line != "controlpath none" }) - ); + let proxy = "openshell ssh-proxy --sandbox demo"; + for command in [ + ssh_forward_command(proxy, "sandbox-id", &ForwardSpec::new(18080)), + ssh_upload_command(proxy), + ] { + let output = Command::new("ssh") + .arg("-F") + .arg(&config) + .arg("-G") + .args(command.get_args()) + .output() + .unwrap(); + assert!( + output.status.success(), + "ssh -G failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + let effective = String::from_utf8(output.stdout).unwrap(); + assert!(effective.lines().any(|line| line == "controlmaster false")); + assert!(effective.lines().any(|line| line == "controlpersist no")); + assert!( + effective + .lines() + .any(|line| line == "forkafterauthentication no") + ); + assert!( + !effective + .lines() + .any(|line| { line.starts_with("controlpath ") && line != "controlpath none" }) + ); + } } #[tokio::test] @@ -3551,7 +3563,7 @@ mod tests { let output = run_upload_shell(&command, tmp.path(), Some(&archive)); assert_eq!( output.status.code(), - Some(UPLOAD_DESTINATION_CHANGED_EXIT_CODE) + Some(openshell_core::upload::DESTINATION_CHANGED_EXIT_CODE) ); assert!(String::from_utf8_lossy(&output.stderr).contains("changed type")); if initial_kind == UploadDestinationKind::Directory { diff --git a/crates/openshell-cli/src/ssh/upload.rs b/crates/openshell-cli/src/ssh/upload.rs index 8f3453ee7c..05f234b9da 100644 --- a/crates/openshell-cli/src/ssh/upload.rs +++ b/crates/openshell-cli/src/ssh/upload.rs @@ -8,16 +8,20 @@ use std::process::{Command, ExitStatus, Stdio}; use std::time::Duration; use miette::Diagnostic; -use openshell_core::upload::{DESTINATION_CHANGED_EXIT_CODE, UploadPlan, WINDOW_BYTES}; +use openshell_core::upload::{ + ACK_INTERVAL_BYTES, DESTINATION_CHANGED_EXIT_CODE, UploadPlan, WINDOW_BYTES, +}; use thiserror::Error; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use tokio::process::{Child, Command as TokioCommand}; use tokio::sync::{mpsc, watch}; -use tokio::task::AbortHandle; +use tokio::task::{AbortHandle, JoinHandle}; const NEGOTIATION_TIMEOUT: Duration = Duration::from_secs(10); const CHILD_EXIT_TIMEOUT: Duration = Duration::from_secs(5); const CHUNK_BYTES: usize = 64 * 1024; +// A larger chunk could wait forever for credit the supervisor never reports. +const _: () = assert!(CHUNK_BYTES as u64 <= WINDOW_BYTES - ACK_INTERVAL_BYTES); const MAX_DIAGNOSTIC_BYTES: usize = 16 * 1024; /// The producer can hold at most two chunks in addition to its current buffer /// and the chunk held by the sender. It cannot build an entire archive in RAM. @@ -81,13 +85,11 @@ impl Drop for ChildGuard { } } -struct AbortTasks(Vec); +struct AbortOnDrop(AbortHandle); -impl Drop for AbortTasks { +impl Drop for AbortOnDrop { fn drop(&mut self) { - for task in &self.0 { - task.abort(); - } + self.0.abort(); } } @@ -124,25 +126,33 @@ where .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::piped()); - let child = TokioCommand::from(command) + let mut child = TokioCommand::from(command) .kill_on_drop(true) .spawn() .map_err(|error| UploadError::Transport(error.to_string()))?; + // EOF marks the end of the archive, so a failed or cancelled transfer must + // stop SSH before its stdin closes. The transfer closes stdin only on + // success, and declaring stdin first drops the guard, killing SSH, first. + let mut stdin = Some(child.stdin.take().expect("piped SSH stdin")); + let stdout = child.stdout.take().expect("piped SSH stdout"); + let stderr = child.stderr.take().expect("piped SSH stderr"); let mut child = ChildGuard(Some(child)); - // EOF marks the end of the archive, so a failed transfer must stop SSH - // before its stdin closes. The transfer closes stdin only on success. - let mut stdin = Some(child.child().stdin.take().expect("piped SSH stdin")); - let stdout = child.child().stdout.take().expect("piped SSH stdout"); - let stderr = child.child().stderr.take().expect("piped SSH stderr"); let diagnostics = tokio::spawn(read_diagnostic(stderr)); - let diagnostic_guard = AbortTasks(vec![diagnostics.abort_handle()]); + let _diagnostics = AbortOnDrop(diagnostics.abort_handle()); let transfer = match transfer(stdout, &mut stdin, plan, producer).await { Ok(sent) => Ok(sent), Err(Stopped::Stream(message)) => Err(message), Err(Stopped::Local(error)) => { let _ = child.finish(true).await; - return Err(error); + drop(stdin); + return Err(match error { + UploadError::Transport(message) => UploadError::Transport(format!( + "{message}: {}", + finish_diagnostic(diagnostics).await + )), + error => error, + }); } }; drop(stdin); @@ -153,13 +163,7 @@ where )); }; let status = status.map_err(|error| UploadError::Transport(error.to_string()))?; - let diagnostic = tokio::time::timeout(CHILD_EXIT_TIMEOUT, diagnostics) - .await - .ok() - .and_then(Result::ok) - .unwrap_or_default(); - drop(diagnostic_guard); - let diagnostic = diagnostic.trim(); + let diagnostic = finish_diagnostic(diagnostics).await; match status.code() { Some(0) => match transfer { Ok(sent) if sent.acknowledged == sent.bytes => Ok(()), @@ -191,6 +195,17 @@ where } } +/// The tail of SSH stderr, which includes extraction output. +async fn finish_diagnostic(task: JoinHandle) -> String { + tokio::time::timeout(CHILD_EXIT_TIMEOUT, task) + .await + .ok() + .and_then(Result::ok) + .unwrap_or_default() + .trim() + .to_string() +} + async fn read_diagnostic(mut reader: R) -> String { let mut diagnostic = Vec::new(); let mut chunk = [0; 4096]; @@ -316,7 +331,7 @@ where } Ok(()) }); - let _acks = AbortTasks(vec![acks.abort_handle()]); + let _acks = AbortOnDrop(acks.abort_handle()); let mut chunks = spawn_producer(producer); let mut sent = 0_u64; loop { @@ -350,8 +365,6 @@ where }); } } - // Count bytes before writing: the supervisor may acknowledge part of - // a chunk while the pipe write is still in progress. sent += length; writer.write_all(&chunk).await.map_err(stream)?; } @@ -380,7 +393,8 @@ mod tests { use std::sync::atomic::{AtomicBool, Ordering}; use tokio::io::{AsyncBufReadExt, BufReader, DuplexStream, ReadHalf, WriteHalf}; - const WINDOW: usize = 1024 * 1024; + #[allow(clippy::cast_possible_truncation)] + const WINDOW: usize = WINDOW_BYTES as usize; type Check = fn(&Result<(), UploadError>) -> bool; fn plan() -> UploadPlan { @@ -393,7 +407,7 @@ mod tests { fn start( producer: impl FnOnce(&mut dyn Write) -> miette::Result<()> + Send + 'static, buffer: usize, - ) -> (tokio::task::JoinHandle>, DuplexStream) { + ) -> (JoinHandle>, DuplexStream) { let (client, server) = tokio::io::duplex(buffer); let (read, write) = tokio::io::split(client); let plan = plan().encode().unwrap(); @@ -592,7 +606,6 @@ mod tests { // archive before reporting. `\0\0\0\0\0\0\0\a` is a total of seven bytes. const ZERO: &str = r"printf '\0\0\0\0\0\0\0\0'"; const SEVEN: &str = r"printf '\0\0\0\0\0\0\0\a'"; - assert_eq!(WINDOW as u64, WINDOW_BYTES); let cases: [(String, Check); 4] = [ ( format!("read -r plan; {ZERO}; cat >/dev/null; {SEVEN}"), diff --git a/crates/openshell-core/src/upload.rs b/crates/openshell-core/src/upload.rs index f3939a5aa4..1483e1fe12 100644 --- a/crates/openshell-core/src/upload.rs +++ b/crates/openshell-core/src/upload.rs @@ -18,6 +18,12 @@ pub const SUBSYSTEM: &str = "openshell-upload-v1"; /// Archive bytes a client may send beyond the last acknowledged total. This /// must stay below the supervisor's pending-stdin budget. pub const WINDOW_BYTES: u64 = 1024 * 1024; +/// Archive bytes between progress reports; success also reports the final total. +/// +/// Clients must write archive chunks of at most +/// `WINDOW_BYTES - ACK_INTERVAL_BYTES`, so a drained extractor always returns +/// enough credit for the next chunk. +pub const ACK_INTERVAL_BYTES: u64 = 256 * 1024; pub const MAX_PLAN_BYTES: usize = 16 * 1024; pub const DESTINATION_CHANGED_EXIT_CODE: i32 = 73; diff --git a/crates/openshell-supervisor-process/README.md b/crates/openshell-supervisor-process/README.md index 7f7b43040a..1c2f728b43 100644 --- a/crates/openshell-supervisor-process/README.md +++ b/crates/openshell-supervisor-process/README.md @@ -27,8 +27,8 @@ After the extractor exits successfully and stdin has drained, the supervisor sends the final total, then the exit status. The client succeeds only when the status is zero and the final total equals the bytes it sent. A malformed plan, credit violation, extended data or EOF before the plan fails the upload with -exit status 74. Whichever happens first, the exit or the rejection, decides the -reported result. Rejection, channel close or disconnect terminates the +exit status 74. A rejection wins until the supervisor reports the exit result, +even after the extractor exits. Rejection, channel close or disconnect terminates the extractor. Extraction writes directly to the destination, so failures can leave partially extracted files; there is no rollback or atomic replacement. diff --git a/crates/openshell-supervisor-process/src/ssh.rs b/crates/openshell-supervisor-process/src/ssh.rs index 7e2601c34f..905a5f0a38 100644 --- a/crates/openshell-supervisor-process/src/ssh.rs +++ b/crates/openshell-supervisor-process/src/ssh.rs @@ -1181,7 +1181,7 @@ impl SshHandler { if let Some(message) = message { session.extended_data(channel, 1, format!("openshell: {message}\n").into_bytes())?; } - self.reject_exec_input(channel, session, "invalid upload stream") + self.reject_exec_input(channel, session, message.unwrap_or("invalid upload stream")) } fn reject_exec_input( @@ -1197,7 +1197,7 @@ impl SshHandler { if let Some(process) = state.process.take() { tokio::spawn(async move { if let Err(error) = process.terminate().await { - warn!(%error, "failed to terminate exec after stdin overflow"); + warn!(%error, "failed to terminate rejected exec"); } }); } diff --git a/crates/openshell-supervisor-process/src/ssh/upload.rs b/crates/openshell-supervisor-process/src/ssh/upload.rs index 11b3363719..c8656f8705 100644 --- a/crates/openshell-supervisor-process/src/ssh/upload.rs +++ b/crates/openshell-supervisor-process/src/ssh/upload.rs @@ -9,7 +9,7 @@ //! credit. The SSH exit status reports the result. use super::input::{InputProgress, InputState}; -use openshell_core::upload::{PlanError, PlanReader, UploadPlan, WINDOW_BYTES}; +use openshell_core::upload::{ACK_INTERVAL_BYTES, PlanError, PlanReader, UploadPlan, WINDOW_BYTES}; use openshell_isolation_interface::contract::BoundaryProcess; use russh::ChannelId; use russh::server::Handle; @@ -19,7 +19,8 @@ use std::time::Duration; use tokio::sync::watch; use tokio::task::{AbortHandle, JoinHandle}; -const ACK_INTERVAL: u64 = 256 * 1024; +// Honest clients cannot reach the general pending-input guard. +const _: () = assert!(WINDOW_BYTES < super::input::MAX_PENDING_INPUT as u64); /// A successful `cat | tar` exit follows stdin EOF, so the writer publishes /// its drained state almost immediately. const DRAIN_TIMEOUT: Duration = Duration::from_secs(5); @@ -181,7 +182,7 @@ async fn send_acks( let _ = process.terminate().await; return; } - if latest.consumed - acknowledged >= ACK_INTERVAL { + if latest.consumed - acknowledged >= ACK_INTERVAL_BYTES { if handle.data(channel, ack(latest.consumed)).await.is_err() { return; } diff --git a/crates/openshell-supervisor-process/src/ssh/upload_tests.rs b/crates/openshell-supervisor-process/src/ssh/upload_tests.rs index b49cb6bf77..196dccac79 100644 --- a/crates/openshell-supervisor-process/src/ssh/upload_tests.rs +++ b/crates/openshell-supervisor-process/src/ssh/upload_tests.rs @@ -18,7 +18,8 @@ use tokio::sync::{mpsc, oneshot, watch}; const DEADLINE: Duration = Duration::from_secs(10); const CHUNK: usize = 64 * 1024; -const WINDOW: usize = 1024 * 1024; +#[allow(clippy::cast_possible_truncation)] +const WINDOW: usize = WINDOW_BYTES as usize; struct Process { status: watch::Sender>, @@ -192,7 +193,6 @@ async fn start(client: &russh::client::Handle) -> (Channel, #[tokio::test] async fn stalled_upload_preserves_controls_then_resumes_beyond_four_mib() { - assert_eq!(WINDOW as u64, WINDOW_BYTES); tokio::time::timeout(DEADLINE, async { let (upload, mut control) = exec_session(); let (sibling, mut other) = exec_session(); @@ -266,6 +266,34 @@ async fn stalled_upload_preserves_controls_then_resumes_beyond_four_mib() { .expect("upload did not resume or blocked SSH control traffic"); } +#[tokio::test] +async fn small_upload_in_one_packet_reports_final_total() { + tokio::time::timeout(DEADLINE, async { + let (session, mut control) = exec_session(); + let client = client_for(vec![session]).await; + let mut channel = open(&client).await; + // The plan and archive share one packet; the tail stays below the + // progress interval, so only the final total covers it. + let mut packet = plan_line(); + packet.push(b'x'); + channel.data(packet.as_slice()).await.unwrap(); + channel.eof().await.unwrap(); + let mut bytes = Vec::new(); + control.stdin.read_to_end(&mut bytes).await.unwrap(); + assert_eq!(bytes, b"x"); + control + .process + .status + .send_replace(Some(BoundaryExitStatus::Exited(0))); + drop((control.stdout, control.stderr)); + let mut reader = Reader::default(); + assert_eq!(reader.close(&mut channel).await, [0, 1]); + assert_eq!(reader.statuses, [0]); + }) + .await + .expect("small upload did not report its final total"); +} + #[tokio::test] async fn invalid_plans_and_credit_violations_exit_74() { for case in ["credit", "eof", "plan"] { From d5bbc71fd83175473dddcbd207a838455b74a6d2 Mon Sep 17 00:00:00 2001 From: Shiju Date: Tue, 13 Oct 2026 13:04:29 +0530 Subject: [PATCH 5/5] test(cli): assert upload ssh opens the subsystem without a TTY Signed-off-by: Shiju --- crates/openshell-cli/src/ssh.rs | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/crates/openshell-cli/src/ssh.rs b/crates/openshell-cli/src/ssh.rs index 420606e70b..a576d5a788 100644 --- a/crates/openshell-cli/src/ssh.rs +++ b/crates/openshell-cli/src/ssh.rs @@ -2924,11 +2924,13 @@ mod tests { let upload = ssh_upload_command("openshell ssh-proxy --sandbox demo"); let effective = effective_ssh_config(&upload, &[]); assert_connection_sharing_disabled(&effective); - assert!( - effective - .lines() - .any(|line| line == "forkafterauthentication no") - ); + for expected in [ + "forkafterauthentication no", + "requesttty false", + "sessiontype subsystem", + ] { + assert!(effective.lines().any(|line| line == expected), "{expected}"); + } } #[tokio::test]