diff --git a/src/boxlite/src/portal/interfaces/exec.rs b/src/boxlite/src/portal/interfaces/exec.rs index b04135b9d..f87771d0e 100644 --- a/src/boxlite/src/portal/interfaces/exec.rs +++ b/src/boxlite/src/portal/interfaces/exec.rs @@ -337,8 +337,8 @@ impl ExecProtocol { // cut, doubling visible columns and desyncing TUI cursor // math (see https://github.com/.../issues/...). Holding // the trailing partial across chunks fixes this. - let mut stdout = DecodedStream::new(stdout_tx); - let mut stderr = DecodedStream::new(stderr_tx); + let mut stdout = OutputTracker::new(stdout_tx); + let mut stderr = OutputTracker::new(stderr_tx); loop { // Use select! to handle cancellation while streaming @@ -375,7 +375,8 @@ impl ExecProtocol { // synthesized "Attach stream error: …" line. stdout.flush(); stderr.flush(); - let _ = stderr.tx.send(format!("Attach stream error: {}", e)); + let _ = + stderr.stream.tx.send(format!("Attach stream error: {}", e)); break; } None => { @@ -404,20 +405,66 @@ impl ExecProtocol { }); } - fn route_output(output: ExecOutput, stdout: &mut DecodedStream, stderr: &mut DecodedStream) { + fn route_output(output: ExecOutput, stdout: &mut OutputTracker, stderr: &mut OutputTracker) { match output.event { Some(exec_output::Event::Stdout(chunk)) => { tracing::trace!(len = chunk.data.len(), "Received exec stdout"); - stdout.send_bytes(chunk.data); + Self::forward_output( + stdout, + "stdout", + chunk.offset, + chunk.data, + chunk.total_bytes, + ); } Some(exec_output::Event::Stderr(chunk)) => { tracing::trace!(len = chunk.data.len(), "Received exec stderr"); - stderr.send_bytes(chunk.data); + Self::forward_output( + stderr, + "stderr", + chunk.offset, + chunk.data, + chunk.total_bytes, + ); } None => {} } } + fn forward_output( + output: &mut OutputTracker, + source: &str, + offset: Option, + data: Vec, + total_bytes: Option, + ) { + match output.receive(offset, data, total_bytes) { + ReceivedOutput::Data { data, lost_bytes } => { + if let Some(lost_bytes) = lost_bytes { + Self::report_gap(output, source, lost_bytes); + } + output.stream.send_bytes(data); + } + ReceivedOutput::End { lost_bytes } => { + if let Some(lost_bytes) = lost_bytes { + Self::report_gap(output, source, lost_bytes); + } + } + ReceivedOutput::Ignore => {} + } + } + + fn report_gap(output: &mut OutputTracker, source: &str, lost_bytes: u64) { + tracing::warn!( + source, + lost_bytes, + "Guest output buffer dropped older output" + ); + let _ = output.stream.tx.send(format!( + "[boxlite] {source} output dropped {lost_bytes} bytes\r\n" + )); + } + fn spawn_wait( mut client: ExecutionClient, execution_id: String, @@ -719,6 +766,95 @@ impl DecodedStream { } } +struct OutputTracker { + expected_offset: u64, + stream: DecodedStream, +} + +enum ReceivedOutput { + Data { + data: Vec, + lost_bytes: Option, + }, + End { + lost_bytes: Option, + }, + Ignore, +} + +impl OutputTracker { + fn new(tx: mpsc::UnboundedSender) -> Self { + Self { + expected_offset: 0, + stream: DecodedStream::new(tx), + } + } + + fn receive( + &mut self, + offset: Option, + data: Vec, + total_bytes: Option, + ) -> ReceivedOutput { + if let Some(total_bytes) = total_bytes { + let Some(offset) = offset else { + tracing::warn!(total_bytes, "Exec output end frame has no offset"); + return ReceivedOutput::Ignore; + }; + if !data.is_empty() || offset != total_bytes { + tracing::warn!(offset, total_bytes, "Invalid exec output end frame"); + return ReceivedOutput::Ignore; + } + if total_bytes < self.expected_offset { + tracing::warn!( + expected_offset = self.expected_offset, + total_bytes, + "Exec output end frame moved backwards" + ); + return ReceivedOutput::Ignore; + } + + let lost_bytes = total_bytes - self.expected_offset; + self.expected_offset = total_bytes; + self.stream.flush(); + return ReceivedOutput::End { + lost_bytes: (lost_bytes > 0).then_some(lost_bytes), + }; + } + + let Some(offset) = offset else { + self.expected_offset += data.len() as u64; + return ReceivedOutput::Data { + data, + lost_bytes: None, + }; + }; + + if offset < self.expected_offset { + tracing::warn!( + expected_offset = self.expected_offset, + offset, + "Exec output chunk moved backwards" + ); + return ReceivedOutput::Ignore; + } + + let lost_bytes = offset - self.expected_offset; + if lost_bytes > 0 { + self.stream.flush(); + } + self.expected_offset = offset + data.len() as u64; + ReceivedOutput::Data { + data, + lost_bytes: (lost_bytes > 0).then_some(lost_bytes), + } + } + + fn flush(&mut self) { + self.stream.flush(); + } +} + // ============================================================================ // UNIT TESTS // ============================================================================ @@ -1146,16 +1282,20 @@ mod tests { let (stdout_tx, mut stdout_rx) = mpsc::unbounded_channel::(); let (stderr_tx, mut stderr_rx) = mpsc::unbounded_channel::(); - let mut stdout = DecodedStream::new(stdout_tx); - let mut stderr = DecodedStream::new(stderr_tx); - - let mk_stdout = |bytes: Vec| ExecOutput { - event: Some(exec_output::Event::Stdout(StdoutMsg { data: bytes })), + let mut stdout = OutputTracker::new(stdout_tx); + let mut stderr = OutputTracker::new(stderr_tx); + + let mk_stdout = |offset, bytes: Vec| ExecOutput { + event: Some(exec_output::Event::Stdout(StdoutMsg { + data: bytes, + offset: Some(offset), + total_bytes: None, + })), }; // "─" split into [E2] and [94 80] across two messages. - ExecProtocol::route_output(mk_stdout(vec![0xE2]), &mut stdout, &mut stderr); - ExecProtocol::route_output(mk_stdout(vec![0x94, 0x80]), &mut stdout, &mut stderr); + ExecProtocol::route_output(mk_stdout(0, vec![0xE2]), &mut stdout, &mut stderr); + ExecProtocol::route_output(mk_stdout(1, vec![0x94, 0x80]), &mut stdout, &mut stderr); // First message: holdover only, no emission. // Second message: complete "─" emitted. @@ -1167,6 +1307,113 @@ mod tests { assert!(stderr_rx.try_recv().is_err()); } + #[test] + fn output_gap_only_flushes_the_affected_utf8_decoder() { + use boxlite_shared::{Stderr as StderrMsg, Stdout as StdoutMsg, exec_output}; + + let (stdout_tx, mut stdout_rx) = mpsc::unbounded_channel::(); + let (stderr_tx, mut stderr_rx) = mpsc::unbounded_channel::(); + let mut stdout = OutputTracker::new(stdout_tx); + let mut stderr = OutputTracker::new(stderr_tx); + + let output = |event| ExecOutput { event: Some(event) }; + ExecProtocol::route_output( + output(exec_output::Event::Stderr(StderrMsg { + data: vec![0xE2], + offset: Some(0), + total_bytes: None, + })), + &mut stdout, + &mut stderr, + ); + ExecProtocol::route_output( + output(exec_output::Event::Stdout(StdoutMsg { + data: b"after-gap".to_vec(), + offset: Some(1), + total_bytes: None, + })), + &mut stdout, + &mut stderr, + ); + ExecProtocol::route_output( + output(exec_output::Event::Stderr(StderrMsg { + data: vec![0x94, 0x80], + offset: Some(1), + total_bytes: None, + })), + &mut stdout, + &mut stderr, + ); + + assert_eq!( + stdout_rx.try_recv().ok(), + Some("[boxlite] stdout output dropped 1 bytes\r\n".to_string()) + ); + assert_eq!(stdout_rx.try_recv().ok(), Some("after-gap".to_string())); + assert_eq!(stderr_rx.try_recv().ok(), Some("─".to_string())); + } + + #[test] + fn output_end_reports_a_missing_tail() { + use boxlite_shared::{Stdout as StdoutMsg, exec_output}; + + let (stdout_tx, mut stdout_rx) = mpsc::unbounded_channel::(); + let (stderr_tx, mut stderr_rx) = mpsc::unbounded_channel::(); + let mut stdout = OutputTracker::new(stdout_tx); + let mut stderr = OutputTracker::new(stderr_tx); + + let output = |event| ExecOutput { event: Some(event) }; + ExecProtocol::route_output( + output(exec_output::Event::Stdout(StdoutMsg { + data: b"hello".to_vec(), + offset: Some(0), + total_bytes: None, + })), + &mut stdout, + &mut stderr, + ); + ExecProtocol::route_output( + output(exec_output::Event::Stdout(StdoutMsg { + data: Vec::new(), + offset: Some(10), + total_bytes: Some(10), + })), + &mut stdout, + &mut stderr, + ); + + assert_eq!(stdout_rx.try_recv().ok(), Some("hello".to_string())); + assert_eq!( + stdout_rx.try_recv().ok(), + Some("[boxlite] stdout output dropped 5 bytes\r\n".to_string()) + ); + assert!(stderr_rx.try_recv().is_err()); + } + + #[test] + fn legacy_output_without_offsets_remains_contiguous() { + use boxlite_shared::{Stdout as StdoutMsg, exec_output}; + + let (stdout_tx, mut stdout_rx) = mpsc::unbounded_channel::(); + let (stderr_tx, mut stderr_rx) = mpsc::unbounded_channel::(); + let mut stdout = OutputTracker::new(stdout_tx); + let mut stderr = OutputTracker::new(stderr_tx); + + let output = |data| ExecOutput { + event: Some(exec_output::Event::Stdout(StdoutMsg { + data, + offset: None, + total_bytes: None, + })), + }; + ExecProtocol::route_output(output(b"hello ".to_vec()), &mut stdout, &mut stderr); + ExecProtocol::route_output(output(b"world".to_vec()), &mut stdout, &mut stderr); + + assert_eq!(stdout_rx.try_recv().ok(), Some("hello ".to_string())); + assert_eq!(stdout_rx.try_recv().ok(), Some("world".to_string())); + assert!(stderr_rx.try_recv().is_err()); + } + /// Flushing a DecodedStream must drain held-over bytes, leave the /// decoder in a valid drained state, and be idempotent. The attach loop /// flushes both streams on every exit path (clean EOF, transport error, diff --git a/src/boxlite/tests/run_main_command.rs b/src/boxlite/tests/run_main_command.rs index 2654c8706..44e49a64d 100644 --- a/src/boxlite/tests/run_main_command.rs +++ b/src/boxlite/tests/run_main_command.rs @@ -8,7 +8,7 @@ mod common; -use boxlite::{BoxOptions, RootfsSpec}; +use boxlite::{BoxCommand, BoxOptions, LiteBox, RootfsSpec}; use tokio_stream::StreamExt; /// Create a box whose main command is `cmd`, optionally on a PTY. @@ -63,6 +63,122 @@ async fn attached_stdout(opts: BoxOptions) -> String { stdout } +async fn wait_for_file(handle: &LiteBox, path: &str) { + tokio::time::timeout(std::time::Duration::from_secs(30), async { + loop { + let execution = handle + .exec(BoxCommand::new("test").args(["-e", path])) + .await + .expect("check main command marker"); + if execution + .wait() + .await + .expect("wait for marker check") + .exit_code + == 0 + { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + }) + .await + .expect("main command must reach marker"); +} + +#[tokio::test] +async fn main_command_exits_after_large_output_without_attach() { + let home = boxlite_test_utils::home::PerTestBoxHome::new(); + let runtime = boxlite::BoxliteRuntime::new(boxlite::runtime::options::BoxliteOptions { + home_dir: home.path.clone(), + image_registries: common::test_registries(), + }) + .expect("create runtime"); + + let handle = runtime + .create( + main_command_opts(&["sh", "-c", "head -c 1048576 /dev/zero; exit 23"], false), + None, + ) + .await + .expect("create box"); + + let completed = tokio::time::timeout(std::time::Duration::from_secs(30), async { + handle.start().await.expect("start box"); + loop { + if handle.info().await.expect("get box info").status == boxlite::BoxStatus::Stopped { + return true; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + } + }) + .await + .unwrap_or(false); + + let _ = handle.stop().await; + let _ = runtime.remove(handle.id().as_str(), true).await; + let _ = runtime.shutdown(Some(common::TEST_SHUTDOWN_TIMEOUT)).await; + + assert!( + completed, + "a main command that has no Attach consumer must still drain and exit" + ); +} + +#[tokio::test] +async fn late_attach_reports_output_gap() { + let home = boxlite_test_utils::home::PerTestBoxHome::new(); + let runtime = boxlite::BoxliteRuntime::new(boxlite::runtime::options::BoxliteOptions { + home_dir: home.path.clone(), + image_registries: common::test_registries(), + }) + .expect("create runtime"); + + let handle = runtime + .create( + main_command_opts( + &[ + "sh", + "-c", + "head -c 2097152 /dev/zero; touch /tmp/main-output-ready; sleep 30", + ], + false, + ), + None, + ) + .await + .expect("create box"); + + handle.start().await.expect("start box"); + wait_for_file(&handle, "/tmp/main-output-ready").await; + + let mut execution = handle + .attach(None) + .await + .expect("attach to the main command"); + let mut stdout = execution.stdout().expect("stdout stream"); + let dropped = tokio::time::timeout(std::time::Duration::from_secs(5), async { + while let Some(chunk) = stdout.next().await { + if chunk.contains("[boxlite] stdout output dropped") { + return Some(chunk); + } + } + None + }) + .await + .unwrap_or(None); + + let _ = handle.stop().await; + let _ = runtime.remove(handle.id().as_str(), true).await; + let _ = runtime.shutdown(Some(common::TEST_SHUTDOWN_TIMEOUT)).await; + + let dropped = dropped.expect("late attach must report overwritten output"); + assert!( + dropped.contains("stdout output dropped"), + "late attach must report the stdout gap: {dropped:?}" + ); +} + /// `tty: true` must give the *main command* a real terminal. /// /// `test -t 0` asks the kernel, so this cannot pass unless init's fd 0 really diff --git a/src/deps/libkrun-sys/build.rs b/src/deps/libkrun-sys/build.rs index 925a23d55..e1b8bd797 100644 --- a/src/deps/libkrun-sys/build.rs +++ b/src/deps/libkrun-sys/build.rs @@ -476,7 +476,10 @@ impl LibFixup { let mut cmd = Command::new("patchelf"); cmd.args(["--add-needed", LIBC, lib_path_str]); run_command(&mut cmd, &format!("add {} dependency", LIBC)); - println!("cargo:warning=Added {} dependency to {}", LIBC, lib_path_str); + println!( + "cargo:warning=Added {} dependency to {}", + LIBC, lib_path_str + ); } /// Extract SONAME from versioned library filename. diff --git a/src/guest/src/container/command.rs b/src/guest/src/container/command.rs index 7cd2b62a7..c3e9cd63c 100644 --- a/src/guest/src/container/command.rs +++ b/src/guest/src/container/command.rs @@ -307,12 +307,13 @@ impl ContainerCommand { tracing::debug!(pid = pid.as_raw(), "spawned with pipes"); // Non-PTY mode: stdout and stderr are separate pipes - Ok(ExecHandle::new( - pid, - stdin_write, - stdout_read, - Some(stderr_read), - )) + match ExecHandle::new(pid, stdin_write, stdout_read, Some(stderr_read)) { + Ok(handle) => Ok(handle), + Err(error) => { + terminate_process(pid); + Err(error) + } + } } /// Build phase of PTY spawn: zygote IPC only, no console-socket handshake. @@ -493,7 +494,13 @@ impl SpawnedPty { /// Call this AFTER releasing the container mutex — the handshake /// does not need serialization. pub(crate) fn finish(self) -> BoxliteResult { - let pty_master = self.socket.receive_pty_master()?; + let pty_master = match self.socket.receive_pty_master() { + Ok(pty_master) => pty_master, + Err(error) => { + terminate_process(self.pid); + return Err(error); + } + }; create_pty_child(self.pid, pty_master, self.config) } } @@ -510,18 +517,42 @@ pub(crate) fn create_pty_child( pty_master: OwnedFd, config: PtyConfig, ) -> BoxliteResult { - set_pty_window_size(&pty_master, &config)?; - crate::service::exec::tty::apply_modes(&pty_master, &config.modes)?; - let (stdin, stdout) = reconcile_pty_fds(&pty_master)?; + if let Err(error) = set_pty_window_size(&pty_master, &config) { + terminate_process(pid); + return Err(error); + } + if let Err(error) = crate::service::exec::tty::apply_modes(&pty_master, &config.modes) { + terminate_process(pid); + return Err(error); + } + let (stdin, stdout) = match reconcile_pty_fds(&pty_master) { + Ok(fds) => fds, + Err(error) => { + terminate_process(pid); + return Err(error); + } + }; // PTY mode: stderr is None (merged into stdout) - let mut child = ExecHandle::new(pid, stdin, stdout, None); + let mut child = match ExecHandle::new(pid, stdin, stdout, None) { + Ok(handle) => handle, + Err(error) => { + terminate_process(pid); + return Err(error); + } + }; let pty_controller = pty_master_to_file(pty_master); child.set_pty(pty_controller, config); Ok(child) } +pub(super) fn terminate_process(pid: Pid) { + let _fence = crate::reaper::reap_fence(); + let _ = nix::sys::signal::kill(pid, nix::sys::signal::Signal::SIGKILL); + let _ = nix::sys::wait::waitpid(pid, None); +} + /// Set PTY terminal window size via ioctl. fn set_pty_window_size(pty_master: &OwnedFd, config: &PtyConfig) -> BoxliteResult<()> { use nix::pty::Winsize; diff --git a/src/guest/src/container/lifecycle.rs b/src/guest/src/container/lifecycle.rs index c2a7bb251..cebb580fa 100644 --- a/src/guest/src/container/lifecycle.rs +++ b/src/guest/src/container/lifecycle.rs @@ -339,7 +339,7 @@ impl Container { stdin, stdout, stderr, - } => ExecHandle::new(pid, stdin, stdout, Some(stderr)), + } => init_pipe_handle(pid, stdin, stdout, stderr)?, // Mirrors the tenant PTY path: the master becomes stdin+stdout and // is retained for window-size ioctls, so ResizeTty reaches the main // command exactly as it reaches an exec. @@ -540,6 +540,21 @@ impl Container { } } +fn init_pipe_handle( + pid: nix::unistd::Pid, + stdin: std::os::fd::OwnedFd, + stdout: std::os::fd::OwnedFd, + stderr: std::os::fd::OwnedFd, +) -> BoxliteResult { + match ExecHandle::new(pid, stdin, stdout, Some(stderr)) { + Ok(handle) => Ok(handle), + Err(error) => { + super::command::terminate_process(pid); + Err(error) + } + } +} + // ==================== // Init Health Check // ==================== @@ -580,3 +595,51 @@ impl Drop for Container { tracing::debug!(container_id = %self.id, "Container cleanup complete"); } } + +#[cfg(test)] +mod tests { + use super::*; + use nix::sys::signal::kill; + use nix::unistd::{pipe, Pid}; + use std::os::fd::OwnedFd; + use std::process::{Child, Command}; + + struct ChildGuard(Child); + + impl Drop for ChildGuard { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } + } + + #[test] + fn init_pipe_handle_failure_reaps_init() { + let child = ChildGuard( + Command::new("/bin/sh") + .args(["-c", "sleep 30"]) + .spawn() + .unwrap(), + ); + let pid = Pid::from_raw(child.0.id() as i32); + let (_stdin_read, stdin_write) = pipe().unwrap(); + let stdout: OwnedFd = std::fs::File::open("/proc/self/stat").unwrap().into(); + let (stderr_read, _stderr_write) = pipe().unwrap(); + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_io() + .build() + .unwrap(); + + let result = + runtime.block_on(async { init_pipe_handle(pid, stdin_write, stdout, stderr_read) }); + + assert!( + result.is_err(), + "regular files cannot register with AsyncFd" + ); + assert!( + kill(pid, None).is_err(), + "failed init handle must not leave a child" + ); + } +} diff --git a/src/guest/src/service/exec/exec_handle.rs b/src/guest/src/service/exec/exec_handle.rs index 775f9de2d..edc349615 100644 --- a/src/guest/src/service/exec/exec_handle.rs +++ b/src/guest/src/service/exec/exec_handle.rs @@ -7,26 +7,25 @@ use boxlite_shared::errors::{BoxliteError, BoxliteResult}; use futures::stream::{Stream, StreamExt}; use nix::sys::signal::Signal; use nix::unistd::Pid; -use std::os::unix::io::OwnedFd; +use std::io; +use std::os::fd::{AsRawFd, OwnedFd}; use std::pin::Pin; use std::task::{Context, Poll}; -use tokio::io::AsyncWriteExt; +use tokio::io::{unix::AsyncFd, Interest}; /// Stdin writer for executed process /// /// Async wrapper around file descriptor for writing to process stdin. pub struct ExecStdin { - inner: tokio::fs::File, + inner: AsyncFd, } impl ExecStdin { /// Create from file descriptor - pub fn new(fd: OwnedFd) -> Self { - use std::os::fd::{FromRawFd, IntoRawFd}; - let std_file = unsafe { std::fs::File::from_raw_fd(fd.into_raw_fd()) }; - Self { - inner: tokio::fs::File::from_std(std_file), - } + pub fn new(fd: OwnedFd) -> BoxliteResult { + Ok(Self { + inner: async_fd(fd, "stdin")?, + }) } /// Write all data to stdin @@ -35,49 +34,98 @@ impl ExecStdin { /// /// - I/O error (pipe closed, etc.) pub async fn write_all(&mut self, data: &[u8]) -> BoxliteResult<()> { - self.inner - .write_all(data) - .await - .map_err(|e| BoxliteError::Internal(format!("Failed to write to stdin: {}", e))) + let mut written = 0; + while written < data.len() { + let count = self + .inner + .async_io(Interest::WRITABLE, |fd| write_fd(fd, &data[written..])) + .await + .map_err(|error| { + BoxliteError::Internal(format!("Failed to write to stdin: {error}")) + })?; + if count == 0 { + return Err(BoxliteError::Internal( + "Failed to write to stdin: write returned zero bytes".into(), + )); + } + written += count; + } + Ok(()) + } +} + +fn async_fd(fd: OwnedFd, stream_name: &str) -> BoxliteResult> { + set_nonblocking(&fd)?; + AsyncFd::new(fd).map_err(|error| { + BoxliteError::Internal(format!( + "Failed to register {stream_name} with Tokio I/O reactor: {error}" + )) + }) +} + +fn set_nonblocking(fd: &OwnedFd) -> BoxliteResult<()> { + let flags = nix::fcntl::fcntl(fd.as_raw_fd(), nix::fcntl::FcntlArg::F_GETFL) + .map_err(|error| BoxliteError::Internal(format!("Failed to read fd flags: {error}")))?; + let mut flags = nix::fcntl::OFlag::from_bits_truncate(flags); + flags.insert(nix::fcntl::OFlag::O_NONBLOCK); + nix::fcntl::fcntl(fd.as_raw_fd(), nix::fcntl::FcntlArg::F_SETFL(flags)).map_err(|error| { + BoxliteError::Internal(format!("Failed to set fd non-blocking: {error}")) + })?; + Ok(()) +} + +fn read_fd(fd: &OwnedFd, buffer: &mut [u8]) -> io::Result { + loop { + match nix::unistd::read(fd.as_raw_fd(), buffer) { + Err(nix::errno::Errno::EINTR) => continue, + result => return result.map_err(Into::into), + } + } +} + +fn write_fd(fd: &OwnedFd, buffer: &[u8]) -> io::Result { + loop { + match nix::unistd::write(fd, buffer) { + Err(nix::errno::Errno::EINTR) => continue, + result => return result.map_err(Into::into), + } } } // Shared output stream implementation struct OutputStream { - inner: Pin> + Send>>, + inner: Pin>> + Send>>, } impl OutputStream { - fn new(fd: OwnedFd) -> Self { + fn new(fd: OwnedFd, pty_eio_is_eof: bool) -> BoxliteResult { use async_stream::stream; - use std::os::fd::{FromRawFd, IntoRawFd}; - use tokio::io::AsyncReadExt; - // Convert OwnedFd to tokio file - let std_file = unsafe { std::fs::File::from_raw_fd(fd.into_raw_fd()) }; - let file = tokio::fs::File::from_std(std_file); - let mut reader = tokio::io::BufReader::new(file); + let reader = async_fd(fd, "output")?; - // Read chunks as they arrive (works for both PTY and pipes) let stream = stream! { let mut buf = [0u8; 1024]; loop { - match reader.read(&mut buf).await { + match reader.async_io(Interest::READABLE, |fd| read_fd(fd, &mut buf)).await { Ok(0) => break, // EOF - Ok(n) => yield buf[..n].to_vec(), - Err(_) => break, + Ok(n) => yield Ok(buf[..n].to_vec()), + Err(error) if pty_eio_is_eof && error.raw_os_error() == Some(nix::libc::EIO) => break, + Err(error) => { + yield Err(error); + break; + } } } }; - Self { + Ok(Self { inner: Box::pin(stream), - } + }) } } impl Stream for OutputStream { - type Item = Vec; + type Item = io::Result>; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.inner.as_mut().poll_next(cx) @@ -93,15 +141,21 @@ pub struct ExecStdout { impl ExecStdout { /// Create from file descriptor - pub fn new(fd: OwnedFd) -> Self { - Self { - inner: OutputStream::new(fd), - } + pub fn new(fd: OwnedFd) -> BoxliteResult { + Ok(Self { + inner: OutputStream::new(fd, false)?, + }) + } + + fn new_pty(fd: OwnedFd) -> BoxliteResult { + Ok(Self { + inner: OutputStream::new(fd, true)?, + }) } } impl Stream for ExecStdout { - type Item = Vec; + type Item = io::Result>; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.inner.poll_next_unpin(cx) @@ -117,15 +171,15 @@ pub struct ExecStderr { impl ExecStderr { /// Create from file descriptor - pub fn new(fd: OwnedFd) -> Self { - Self { - inner: OutputStream::new(fd), - } + pub fn new(fd: OwnedFd) -> BoxliteResult { + Ok(Self { + inner: OutputStream::new(fd, false)?, + }) } } impl Stream for ExecStderr { - type Item = Vec; + type Item = io::Result>; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.inner.poll_next_unpin(cx) @@ -265,17 +319,27 @@ impl ExecHandle { /// * `stdin` - Stdin file descriptor /// * `stdout` - Stdout file descriptor /// * `stderr` - Stderr file descriptor, or `None` in PTY mode (merged into stdout) - pub fn new(pid: Pid, stdin: OwnedFd, stdout: OwnedFd, stderr: Option) -> Self { - Self { + pub fn new( + pid: Pid, + stdin: OwnedFd, + stdout: OwnedFd, + stderr: Option, + ) -> BoxliteResult { + let stdout = if stderr.is_some() { + ExecStdout::new(stdout)? + } else { + ExecStdout::new_pty(stdout)? + }; + Ok(Self { pid, - stdin: Some(ExecStdin::new(stdin)), - stdout: Some(ExecStdout::new(stdout)), + stdin: Some(ExecStdin::new(stdin)?), + stdout: Some(stdout), // In PTY mode, stderr is None because stdout/stderr are merged // at the PTY level (single reader from PTY master) - stderr: stderr.map(ExecStderr::new), + stderr: stderr.map(ExecStderr::new).transpose()?, pty_controller: None, pty_config: None, - } + }) } /// Set PTY controller and config @@ -443,8 +507,8 @@ mod process_group_tests { assert!(checked_process_group_target(Pid::from_raw(0), Pid::from_raw(0)).is_err()); } - #[test] - fn process_group_kill_reaches_a_background_descendant() { + #[tokio::test] + async fn process_group_kill_reaches_a_background_descendant() { let mut command = Command::new("/bin/sh"); command .arg("-c") @@ -481,7 +545,8 @@ mod process_group_tests { let (_stdin_peer, stdin) = nix::unistd::pipe().unwrap(); let (stdout, _stdout_peer) = nix::unistd::pipe().unwrap(); let (stderr, _stderr_peer) = nix::unistd::pipe().unwrap(); - let handle = ExecHandle::new(leader, stdin, stdout, Some(stderr)); + let handle = ExecHandle::new(leader, stdin, stdout, Some(stderr)) + .expect("test pipe must register with Tokio"); handle.kill_process_group(Signal::SIGTERM).unwrap(); std::thread::sleep(Duration::from_millis(25)); @@ -504,3 +569,52 @@ mod process_group_tests { ); } } + +#[cfg(test)] +mod tests { + use super::*; + use std::task::Poll; + use std::time::Duration; + + #[test] + fn output_read_does_not_occupy_blocking_pool() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .max_blocking_threads(1) + .enable_all() + .build() + .unwrap(); + runtime.block_on(async { + let (read_fd, _write_fd) = nix::unistd::pipe().unwrap(); + let mut stdout = ExecStdout::new(read_fd).unwrap(); + let (pending_tx, pending_rx) = tokio::sync::oneshot::channel(); + let output_read = tokio::spawn(async move { + let mut pending_tx = Some(pending_tx); + futures::future::poll_fn(move |cx| { + let next = Pin::new(&mut stdout).poll_next(cx); + if matches!(next, Poll::Pending) { + if let Some(tx) = pending_tx.take() { + let _ = tx.send(()); + } + } + next + }) + .await + }); + pending_rx.await.unwrap(); + + let temporary_file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(temporary_file.path(), b"ready").unwrap(); + let read_file = tokio::fs::read(temporary_file.path()); + + assert!( + tokio::time::timeout(Duration::from_millis(100), read_file) + .await + .is_ok(), + "an idle output stream must not occupy a blocking-pool worker" + ); + + output_read.abort(); + }); + } +} diff --git a/src/guest/src/service/exec/executor.rs b/src/guest/src/service/exec/executor.rs index b4640baf9..f87932107 100644 --- a/src/guest/src/service/exec/executor.rs +++ b/src/guest/src/service/exec/executor.rs @@ -161,19 +161,25 @@ fn spawn_with_pipes(req: &ExecRequest) -> BoxliteResult { cmd.stderr(std::process::Stdio::from_raw_fd(stderr_write.into_raw_fd())); } - let child = cmd + let mut child = cmd .spawn() .map_err(|e| BoxliteError::Internal(format!("Failed to spawn '{}': {}", req.program, e)))?; let pid = child.id(); // Non-PTY mode: stdout and stderr are separate pipes - Ok(ExecHandle::new( + match ExecHandle::new( Pid::from_raw(pid as i32), stdin_write, stdout_read, Some(stderr_read), - )) + ) { + Ok(handle) => Ok(handle), + Err(error) => { + terminate_child(&mut child); + Err(error) + } + } } /// Spawn process with PTY (interactive mode). @@ -242,7 +248,7 @@ fn spawn_with_pty(req: &ExecRequest, config: PtyConfig) -> BoxliteResult BoxliteResult fd, + Err(error) => { + terminate_child(&mut child); + return Err(BoxliteError::Internal(format!( + "Failed to dup PTY for stdin: {error}" + ))); + } + }; + let stdout_fd = match dup(master.as_raw_fd()) { + Ok(fd) => fd, + Err(error) => { + terminate_child(&mut child); + return Err(BoxliteError::Internal(format!( + "Failed to dup PTY for stdout: {error}" + ))); + } + }; let stdin = unsafe { OwnedFd::from_raw_fd(stdin_fd) }; let stdout = unsafe { OwnedFd::from_raw_fd(stdout_fd) }; // PTY mode: stderr is None (merged into stdout) - let mut handle = ExecHandle::new(Pid::from_raw(pid as i32), stdin, stdout, None); + let mut handle = match ExecHandle::new(Pid::from_raw(pid as i32), stdin, stdout, None) { + Ok(handle) => handle, + Err(error) => { + terminate_child(&mut child); + return Err(error); + } + }; // Keep master FD for resize operations let pty_controller = { @@ -277,3 +303,26 @@ fn spawn_with_pty(req: &ExecRequest, config: PtyConfig) -> BoxliteResult Result, ExecutionError> { + ) -> Result>, ExecutionError> { info!(execution_id = %exec_id, "attach request"); self.execution(exec_id).await?.attach(exec_id).await } @@ -194,7 +195,7 @@ impl Execution for GuestServer { let rx = self .attach_execution(&request.into_inner().execution_id) .await?; - let stream = ReceiverStream::new(rx).map(Ok); + let stream = ReceiverStream::new(rx); Ok(Response::new(Box::pin(stream) as Self::AttachStream)) } diff --git a/src/guest/src/service/exec/output.rs b/src/guest/src/service/exec/output.rs new file mode 100644 index 000000000..47c2ac870 --- /dev/null +++ b/src/guest/src/service/exec/output.rs @@ -0,0 +1,476 @@ +use crate::service::exec::exec_handle::{ExecStderr, ExecStdout}; +use async_stream::stream; +use boxlite_shared::{exec_output, ExecOutput, Stderr, Stdout}; +use futures::{Stream, StreamExt}; +use std::collections::VecDeque; +use std::pin::Pin; +use std::sync::{Arc, Mutex as StdMutex}; +use tokio::sync::{watch, Mutex}; +use tokio::task::JoinHandle; +use tonic::Status; +use tracing::error; + +const BUFFER_CAPACITY_BYTES: usize = 1024 * 1024; + +pub(crate) type AttachStream = Pin> + Send>>; + +#[derive(Clone)] +pub(crate) struct OutputManager { + inner: Arc>, + updated: watch::Sender<()>, + drain_tasks: Arc>>>, +} + +struct OutputState { + entries: VecDeque, + buffered_bytes: usize, + oldest_sequence: u64, + next_sequence: u64, + failure: Option, + attached: bool, + stdout: StreamState, + stderr: StreamState, +} + +struct ReaderFailure { + sequence: u64, + message: String, +} + +struct OutputEntry { + sequence: u64, + output: ExecOutput, + byte_len: usize, +} + +#[derive(Clone, Copy, Debug)] +enum OutputSource { + Stdout, + Stderr, +} + +struct StreamState { + enabled: bool, + finished: bool, + total_bytes: u64, + last_sequence: Option, +} + +impl OutputManager { + pub(crate) fn new(stdout: Option, stderr: Option) -> Self { + let stdout_enabled = stdout.is_some(); + let stderr_enabled = stderr.is_some(); + let (updated, _) = watch::channel(()); + let manager = Self { + inner: Arc::new(Mutex::new(OutputState { + entries: VecDeque::new(), + buffered_bytes: 0, + oldest_sequence: 0, + next_sequence: 0, + failure: None, + attached: false, + stdout: StreamState { + enabled: stdout_enabled, + finished: !stdout_enabled, + total_bytes: 0, + last_sequence: None, + }, + stderr: StreamState { + enabled: stderr_enabled, + finished: !stderr_enabled, + total_bytes: 0, + last_sequence: None, + }, + })), + updated, + drain_tasks: Arc::new(StdMutex::new(Vec::new())), + }; + + if let Some(stdout) = stdout { + manager.spawn(stdout, OutputSource::Stdout); + } + if let Some(stderr) = stderr { + manager.spawn(stderr, OutputSource::Stderr); + } + + manager + } + + pub(crate) async fn shutdown_drains(&self) { + let tasks = std::mem::take( + &mut *self + .drain_tasks + .lock() + .expect("output drain task lock poisoned"), + ); + for task in &tasks { + task.abort(); + } + for task in tasks { + let _ = task.await; + } + } + + pub(crate) async fn attach(&self) -> Result { + { + let mut state = self.inner.lock().await; + if state.attached { + return Err(Status::already_exists("Already attached")); + } + state.attached = true; + } + + let manager = self.clone(); + let output = stream! { + let mut next_sequence = 0; + let mut stdout_end_sent = false; + let mut stderr_end_sent = false; + let mut updates = manager.updated.subscribe(); + + loop { + enum Next { + Item(ExecOutput), + Error(Status), + Wait, + Done, + } + + let next = { + let state = manager.inner.lock().await; + + if next_sequence < state.oldest_sequence { + next_sequence = state.oldest_sequence; + } + let failure = state + .failure + .as_ref() + .filter(|failure| next_sequence >= failure.sequence); + if let Some(failure) = failure { + Next::Error(Status::internal(failure.message.clone())) + } else if state.stdout.ready_to_end(next_sequence) && !stdout_end_sent { + stdout_end_sent = true; + Next::Item(end_output(OutputSource::Stdout, state.stdout.total_bytes)) + } else if state.stderr.ready_to_end(next_sequence) && !stderr_end_sent { + stderr_end_sent = true; + Next::Item(end_output(OutputSource::Stderr, state.stderr.total_bytes)) + } else if next_sequence < state.next_sequence { + let index = (next_sequence - state.oldest_sequence) as usize; + let output = state + .entries + .get(index) + .expect("ring sequence must exist") + .output + .clone(); + next_sequence += 1; + Next::Item(output) + } else if (!state.stdout.enabled || stdout_end_sent) + && (!state.stderr.enabled || stderr_end_sent) + { + Next::Done + } else { + Next::Wait + } + }; + + match next { + Next::Item(item) => yield Ok(item), + Next::Error(status) => { + yield Err(status); + break; + } + Next::Done => break, + Next::Wait => updates.changed().await.expect("output manager must outlive attach stream"), + } + } + }; + + Ok(Box::pin(output)) + } + + fn spawn(&self, stream: S, source: OutputSource) + where + S: Stream>> + Send + Unpin + 'static, + { + let manager = self.clone(); + let task = tokio::spawn(async move { + manager.drain(stream, source).await; + }); + self.drain_tasks + .lock() + .expect("output drain task lock poisoned") + .push(task); + } + + async fn drain(&self, mut stream: S, source: OutputSource) + where + S: Stream>> + Unpin, + { + while let Some(item) = stream.next().await { + match item { + Ok(data) => self.push(source, data).await, + Err(error) => { + error!(?source, %error, "execution output reader failed"); + self.reader_failed(source, error).await; + return; + } + } + } + self.reader_finished(source).await; + } + + async fn push(&self, source: OutputSource, data: Vec) { + let byte_len = data.len(); + + let mut state = self.inner.lock().await; + let sequence = state.next_sequence; + state.next_sequence += 1; + let stream = state.stream_mut(source); + stream.enabled = true; + let offset = stream.total_bytes; + stream.total_bytes += byte_len as u64; + stream.last_sequence = Some(sequence); + let output = data_output(source, data, offset); + + while state.buffered_bytes + byte_len > BUFFER_CAPACITY_BYTES { + let Some(removed) = state.entries.pop_front() else { + state.oldest_sequence = state.next_sequence; + break; + }; + state.buffered_bytes -= removed.byte_len; + state.oldest_sequence = removed.sequence + 1; + } + + if byte_len <= BUFFER_CAPACITY_BYTES { + state.buffered_bytes += byte_len; + state.entries.push_back(OutputEntry { + sequence, + output, + byte_len, + }); + } + drop(state); + self.updated.send_replace(()); + } + + async fn reader_finished(&self, source: OutputSource) { + let mut state = self.inner.lock().await; + state.stream_mut(source).finished = true; + drop(state); + self.updated.send_replace(()); + } + + async fn reader_failed(&self, source: OutputSource, error: std::io::Error) { + let mut state = self.inner.lock().await; + state.stream_mut(source).finished = true; + if state.failure.is_none() { + state.failure = Some(ReaderFailure { + sequence: state.next_sequence, + message: format!("failed to read {source:?}: {error}"), + }); + } + drop(state); + self.updated.send_replace(()); + } +} + +impl OutputState { + fn stream_mut(&mut self, source: OutputSource) -> &mut StreamState { + match source { + OutputSource::Stdout => &mut self.stdout, + OutputSource::Stderr => &mut self.stderr, + } + } +} + +impl StreamState { + fn ready_to_end(&self, next_sequence: u64) -> bool { + self.enabled + && self.finished + && self + .last_sequence + .is_none_or(|last_sequence| next_sequence > last_sequence) + } +} + +fn data_output(source: OutputSource, data: Vec, offset: u64) -> ExecOutput { + let event = match source { + OutputSource::Stdout => exec_output::Event::Stdout(Stdout { + data, + offset: Some(offset), + total_bytes: None, + }), + OutputSource::Stderr => exec_output::Event::Stderr(Stderr { + data, + offset: Some(offset), + total_bytes: None, + }), + }; + ExecOutput { event: Some(event) } +} + +fn end_output(source: OutputSource, total_bytes: u64) -> ExecOutput { + let event = match source { + OutputSource::Stdout => exec_output::Event::Stdout(Stdout { + data: Vec::new(), + offset: Some(total_bytes), + total_bytes: Some(total_bytes), + }), + OutputSource::Stderr => exec_output::Event::Stderr(Stderr { + data: Vec::new(), + offset: Some(total_bytes), + total_bytes: Some(total_bytes), + }), + }; + ExecOutput { event: Some(event) } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn completed_stream_emits_total_bytes_after_replayed_data() { + let manager = OutputManager::new(None, None); + manager + .push(OutputSource::Stdout, b"already sent".to_vec()) + .await; + + let mut output = manager.attach().await.unwrap(); + let first = output.next().await.unwrap().unwrap(); + let Some(exec_output::Event::Stdout(first)) = first.event else { + panic!("the replayed stdout data must be sent first"); + }; + assert_eq!(first.data, b"already sent"); + assert_eq!(first.offset, Some(0)); + assert_eq!(first.total_bytes, None); + + let end = output.next().await.unwrap().unwrap(); + let Some(exec_output::Event::Stdout(end)) = end.event else { + panic!("the stdout end frame must follow its replayed data"); + }; + assert!(end.data.is_empty()); + assert_eq!(end.offset, Some(b"already sent".len() as u64)); + assert_eq!(end.total_bytes, Some(b"already sent".len() as u64)); + } + + #[tokio::test] + async fn closed_stdout_emits_its_end_before_stderr_closes() { + let (stdout_read, stdout_write) = nix::unistd::pipe().unwrap(); + let (stderr_read, stderr_write) = nix::unistd::pipe().unwrap(); + let manager = OutputManager::new( + Some(ExecStdout::new(stdout_read).unwrap()), + Some(ExecStderr::new(stderr_read).unwrap()), + ); + + nix::unistd::write(&stdout_write, b"out").unwrap(); + drop(stdout_write); + + let mut output = manager.attach().await.unwrap(); + let first = tokio::time::timeout(std::time::Duration::from_secs(1), output.next()) + .await + .expect("stdout data must arrive") + .unwrap() + .unwrap(); + let Some(exec_output::Event::Stdout(first)) = first.event else { + panic!("the first event must be stdout data"); + }; + assert_eq!(first.data, b"out"); + + let end = tokio::time::timeout(std::time::Duration::from_secs(1), output.next()) + .await + .expect("stdout end must not wait for stderr EOF") + .unwrap() + .unwrap(); + let Some(exec_output::Event::Stdout(end)) = end.event else { + panic!("stdout must emit its own end frame"); + }; + assert_eq!(end.total_bytes, Some(3)); + + drop(stderr_write); + } + + #[tokio::test] + async fn end_frames_survive_when_all_data_is_evicted() { + let manager = OutputManager::new(None, None); + manager.push(OutputSource::Stdout, b"lost".to_vec()).await; + manager + .push(OutputSource::Stderr, vec![0; BUFFER_CAPACITY_BYTES + 1]) + .await; + + let mut output = manager.attach().await.unwrap(); + let stdout_end = output.next().await.unwrap().unwrap(); + let Some(exec_output::Event::Stdout(stdout_end)) = stdout_end.event else { + panic!("stdout end frame must survive its evicted data"); + }; + assert!(stdout_end.data.is_empty()); + assert_eq!(stdout_end.total_bytes, Some(4)); + + let stderr_end = output.next().await.unwrap().unwrap(); + let Some(exec_output::Event::Stderr(stderr_end)) = stderr_end.event else { + panic!("stderr end frame must survive its evicted data"); + }; + assert!(stderr_end.data.is_empty()); + assert_eq!( + stderr_end.total_bytes, + Some((BUFFER_CAPACITY_BYTES + 1) as u64) + ); + } + + #[tokio::test] + async fn reader_failure_follows_buffered_output() { + let manager = OutputManager::new(None, None); + manager + .push(OutputSource::Stdout, b"before failure".to_vec()) + .await; + manager + .reader_failed( + OutputSource::Stdout, + std::io::Error::other("simulated pipe failure"), + ) + .await; + + let mut output = manager.attach().await.unwrap(); + let first = output.next().await.unwrap().unwrap(); + let Some(exec_output::Event::Stdout(first)) = first.event else { + panic!("buffered stdout must arrive before the reader failure"); + }; + assert_eq!(first.data, b"before failure"); + + let error = output.next().await.unwrap().unwrap_err(); + assert_eq!(error.code(), tonic::Code::Internal); + assert!(error.message().contains("simulated pipe failure")); + } + + #[tokio::test] + async fn first_reader_failure_stops_output_before_later_reader_failure() { + let manager = OutputManager::new(None, None); + manager + .push(OutputSource::Stdout, b"before failure".to_vec()) + .await; + manager + .reader_failed( + OutputSource::Stdout, + std::io::Error::other("stdout failure"), + ) + .await; + manager + .push(OutputSource::Stderr, b"after failure".to_vec()) + .await; + manager + .reader_failed( + OutputSource::Stderr, + std::io::Error::other("stderr failure"), + ) + .await; + + let mut output = manager.attach().await.unwrap(); + let first = output.next().await.unwrap().unwrap(); + let Some(exec_output::Event::Stdout(first)) = first.event else { + panic!("buffered stdout must arrive before the first reader failure"); + }; + assert_eq!(first.data, b"before failure"); + + let error = output.next().await.unwrap().unwrap_err(); + assert!(error.message().contains("stdout failure")); + } +} diff --git a/src/guest/src/service/exec/registry.rs b/src/guest/src/service/exec/registry.rs index 824661031..0c8b638dd 100644 --- a/src/guest/src/service/exec/registry.rs +++ b/src/guest/src/service/exec/registry.rs @@ -136,7 +136,8 @@ mod release_tests { let (_stdin_peer, stdin) = pipe().unwrap(); let (stdout, _stdout_peer) = pipe().unwrap(); let (stderr, _stderr_peer) = pipe().unwrap(); - let handle = ExecHandle::new(Pid::from_raw(pid), stdin, stdout, Some(stderr)); + let handle = ExecHandle::new(Pid::from_raw(pid), stdin, stdout, Some(stderr)) + .expect("test pipe must register with Tokio"); let exit = ExitSlot::settled_for_test(ExitStatus::Code(7)); if is_init { ExecutionState::new_init_session(handle, exit) diff --git a/src/guest/src/service/exec/state.rs b/src/guest/src/service/exec/state.rs index 36977e68a..36d2b1ff3 100644 --- a/src/guest/src/service/exec/state.rs +++ b/src/guest/src/service/exec/state.rs @@ -1,12 +1,13 @@ use crate::service::exec::error::ExecutionError; use crate::service::exec::exec_handle::ExecHandle; +use crate::service::exec::output::OutputManager; use boxlite_shared::ExecOutput; -use futures::Stream; +use futures::{Stream, StreamExt as _}; use std::os::unix::io::AsRawFd; use std::sync::Arc; use tokio::sync::{mpsc, Mutex}; use tokio::task::{AbortHandle, JoinHandle}; -use tracing::info; +use tonic::Status; /// Abstraction for checking container init health. /// @@ -26,6 +27,7 @@ pub(crate) trait InitHealthCheck: Send + Sync { struct Inner { /// The process handle (owns pid, pty_controller, stdin, stdout, stderr) handle: Option, + output: OutputManager, /// Abort handles for stdin forwarding tasks that may own the taken stdin FD. input_tasks: Vec, /// Stdout/stderr forwarding tasks (set on attach) @@ -54,7 +56,6 @@ pub(crate) struct ExecutionExit { /// Execution state. /// /// Handle owns pid, pty_controller, stdin, stdout, stderr. -/// stdin is taken on send_input(), stdout/stderr are taken on attach(). #[derive(Clone)] pub(crate) struct ExecutionState { inner: Arc>, @@ -65,26 +66,28 @@ pub(crate) struct ExecutionState { } impl ExecutionState { - fn from_inner(inner: Inner, exit: crate::reaper::ExitSlot) -> Self { + fn from_handle( + mut handle: ExecHandle, + init_health: Option>>, + exit: crate::reaper::ExitSlot, + ) -> Self { Self { - inner: Arc::new(Mutex::new(inner)), - exit, - } - } - - /// Create new execution state for a guest-side process. - pub(super) fn new(handle: ExecHandle, exit: crate::reaper::ExitSlot) -> Self { - Self::from_inner( - Inner { + inner: Arc::new(Mutex::new(Inner { + output: OutputManager::new(handle.stdout(), handle.stderr()), handle: Some(handle), input_tasks: Vec::new(), output_tasks: Vec::new(), released: false, timed_out: false, - init_health: None, - }, + init_health, + })), exit, - ) + } + } + + /// Create new execution state for a guest-side process. + pub(super) fn new(handle: ExecHandle, exit: crate::reaper::ExitSlot) -> Self { + Self::from_handle(handle, None, exit) } #[cfg(test)] @@ -104,17 +107,7 @@ impl ExecutionState { init_health: Arc>, exit: crate::reaper::ExitSlot, ) -> Self { - Self::from_inner( - Inner { - handle: Some(handle), - input_tasks: Vec::new(), - output_tasks: Vec::new(), - released: false, - timed_out: false, - init_health: Some(init_health), - }, - exit, - ) + Self::from_handle(handle, Some(init_health), exit) } /// Create execution state for the container's init process itself. @@ -123,17 +116,7 @@ impl ExecutionState { /// reparents to guest main (the boxlite-guest agent process), which owns /// `waitpid(-1)`. See `wait_process`. pub(crate) fn new_init_session(handle: ExecHandle, exit: crate::reaper::ExitSlot) -> Self { - Self::from_inner( - Inner { - handle: Some(handle), - input_tasks: Vec::new(), - output_tasks: Vec::new(), - released: false, - timed_out: false, - init_health: None, - }, - exit, - ) + Self::from_handle(handle, None, exit) } /// Check if the container init process died. @@ -288,88 +271,38 @@ impl ExecutionState { } /// Attach to execution output. - /// - /// Takes stdout/stderr from handle and starts forwarding tasks. - /// Returns stream of output chunks. pub async fn attach( &self, exec_id: &str, - ) -> Result, ExecutionError> { - use boxlite_shared::{exec_output, Stderr, Stdout}; - use futures::StreamExt; - - let (tx, rx) = mpsc::channel(100); - - // Take stdout/stderr from handle - let (stdout, stderr) = { - let mut inner = self.inner.lock().await; - - if !inner.output_tasks.is_empty() { - return Err(ExecutionError::AlreadyAttached); + ) -> Result>, ExecutionError> { + let output = { + let inner = self.inner.lock().await; + if inner.released { + return Err(ExecutionError::HandleUnavailable); } - - let handle = inner - .handle - .as_mut() - .ok_or(ExecutionError::HandleUnavailable)?; - - let stdout = handle.stdout(); - let stderr = handle.stderr(); - - (stdout, stderr) + inner.output.clone() }; - - // Spawn forwarding tasks - let mut tasks = Vec::new(); - - // Spawn stdout forwarding task - let exec_id_string = exec_id.to_string(); - if let Some(mut stdout) = stdout { - let tx = tx.clone(); - let handle = tokio::spawn(async move { - while let Some(chunk) = stdout.next().await { - let msg = ExecOutput { - event: Some(exec_output::Event::Stdout(Stdout { data: chunk })), - }; - if tx.send(msg).await.is_err() { - break; - } - } - info!(execution = ?exec_id_string, "Stdout forwarding task ended"); - }); - tasks.push(handle); - } - - // Spawn stderr forwarding task - let exec_id_string = exec_id.to_string(); - if let Some(mut stderr) = stderr { - let tx = tx.clone(); - let handle = tokio::spawn(async move { - while let Some(chunk) = stderr.next().await { - let msg = ExecOutput { - event: Some(exec_output::Event::Stderr(Stderr { data: chunk })), - }; - if tx.send(msg).await.is_err() { - break; - } - } - info!(execution = ?exec_id_string, "Stderr forwarding task ended"); - }); - tasks.push(handle); - } - - // Store tasks - { - let mut inner = self.inner.lock().await; - if inner.released { - for task in tasks { - task.abort(); + let mut output = output + .attach() + .await + .map_err(|_| ExecutionError::AlreadyAttached)?; + let (tx, rx) = mpsc::channel(100); + let execution_id = exec_id.to_owned(); + let task = tokio::spawn(async move { + while let Some(message) = output.next().await { + if tx.send(message).await.is_err() { + break; } - } else { - inner.output_tasks = tasks; } - } + tracing::info!(%execution_id, "execution output forwarding ended"); + }); + let mut inner = self.inner.lock().await; + if inner.released { + task.abort(); + return Err(ExecutionError::HandleUnavailable); + } + inner.output_tasks.push(task); Ok(rx) } @@ -381,20 +314,27 @@ impl ExecutionState { /// clone still exists. Forwarders are aborted because stdin may already have /// moved out of the handle into one of those tasks. pub(super) async fn release_resources(&self) -> bool { - let mut inner = self.inner.lock().await; - if inner.released { - return false; - } - inner.released = true; + let (output, input_tasks, output_tasks) = { + let mut inner = self.inner.lock().await; + if inner.released { + return false; + } + inner.released = true; + let output = inner.output.clone(); + let input_tasks = std::mem::take(&mut inner.input_tasks); + let output_tasks = std::mem::take(&mut inner.output_tasks); + drop(inner.handle.take()); + drop(inner.init_health.take()); + (output, input_tasks, output_tasks) + }; - for task in inner.input_tasks.drain(..) { + for task in input_tasks { task.abort(); } - for task in inner.output_tasks.drain(..) { + for task in output_tasks { task.abort(); } - drop(inner.handle.take()); - drop(inner.init_health.take()); + output.shutdown_drains().await; true } @@ -477,7 +417,8 @@ mod release_tests { pty_controller.as_raw_fd(), ]; - let mut handle = ExecHandle::new(Pid::from_raw(42_424), stdin, stdout, Some(stderr)); + let mut handle = ExecHandle::new(Pid::from_raw(42_424), stdin, stdout, Some(stderr)) + .expect("test pipe must register with Tokio"); handle.set_pty( std::fs::File::from(pty_controller), PtyConfig { diff --git a/src/guest/src/service/ssh/bridge.rs b/src/guest/src/service/ssh/bridge.rs index 59f2cb228..796abb5b9 100644 --- a/src/guest/src/service/ssh/bridge.rs +++ b/src/guest/src/service/ssh/bridge.rs @@ -198,6 +198,7 @@ impl ChannelBridge { OutputPumpStart::AwaitActivation { output: ready.output, buffered_stdout: ready.buffered_stdout, + stdout_offset: ready.stdout_offset, activate_rx, }, Some(activate_tx), @@ -511,15 +512,17 @@ fn streamlocal_launch(container_id: &str, socket_path: String) -> InternalHelper } struct PreparedStreamlocalOutput { - output: mpsc::Receiver, + output: mpsc::Receiver>, buffered_stdout: Vec, + stdout_offset: u64, } enum OutputPumpStart { Attach, AwaitActivation { - output: mpsc::Receiver, + output: mpsc::Receiver>, buffered_stdout: Vec, + stdout_offset: u64, activate_rx: oneshot::Receiver<()>, }, } @@ -527,10 +530,16 @@ enum OutputPumpStart { #[derive(Default)] struct StreamlocalReadyParser { matched: usize, + stdout_offset: u64, } impl StreamlocalReadyParser { - fn push_stdout(&mut self, data: &[u8]) -> Result>, BridgeError> { + fn push_stdout( + &mut self, + offset: Option, + data: &[u8], + ) -> Result>, BridgeError> { + self.stdout_offset = offset.unwrap_or(self.stdout_offset) + data.len() as u64; let magic = crate::service::ssh::streamlocal::STREAMLOCAL_READY_MAGIC; let prefix_bytes = (magic.len() - self.matched).min(data.len()); if data[..prefix_bytes] != magic[self.matched..self.matched + prefix_bytes] { @@ -559,15 +568,24 @@ async fn await_streamlocal_ready( let mut parser = StreamlocalReadyParser::default(); loop { - let message = output.recv().await.ok_or_else(|| { - BridgeError::Exec("streamlocal helper exited before readiness".into()) - })?; + let message = output + .recv() + .await + .ok_or_else(|| { + BridgeError::Exec("streamlocal helper exited before readiness".into()) + })? + .map_err(|error| { + BridgeError::Exec(format!("streamlocal output failed: {error}")) + })?; match message.event { Some(boxlite_shared::exec_output::Event::Stdout(stdout)) => { - if let Some(buffered_stdout) = parser.push_stdout(&stdout.data)? { + if let Some(buffered_stdout) = + parser.push_stdout(stdout.offset, &stdout.data)? + { return Ok(PreparedStreamlocalOutput { output, buffered_stdout, + stdout_offset: parser.stdout_offset, }); } } @@ -604,7 +622,7 @@ fn spawn_output_pump( mut cancel_rx: oneshot::Receiver<()>, ) -> JoinHandle<()> { tokio::spawn(async move { - let (mut output, buffered_stdout) = match output_start { + let (mut output, buffered_stdout, mut stdout_offset) = match output_start { OutputPumpStart::Attach => { let attached = tokio::select! { _ = &mut cancel_rx => return, @@ -620,7 +638,7 @@ fn spawn_output_pump( Err(_) => Err(ExecutionError::Io("attach timed out".into())), }; match attached { - Ok(output) => (output, Vec::new()), + Ok(output) => (output, Vec::new(), 0), Err(error) => { warn!(%error, %execution_id, "SSH attach failed"); terminate_process_group(server.clone(), execution_id.clone()).await; @@ -638,6 +656,7 @@ fn spawn_output_pump( OutputPumpStart::AwaitActivation { output, buffered_stdout, + stdout_offset, mut activate_rx, } => { let activated = tokio::select! { @@ -647,7 +666,7 @@ fn spawn_output_pump( if activated.is_err() { return; } - (output, buffered_stdout) + (output, buffered_stdout, stdout_offset) } }; @@ -655,6 +674,7 @@ fn spawn_output_pump( let _ = session_handle.data(channel_id, buffered_stdout).await; } + let mut stderr_offset = 0; loop { tokio::select! { // A resolved oneshot receiver must not be polled again by the @@ -664,15 +684,90 @@ fn spawn_output_pump( _ = &mut cancel_rx => return, message = output.recv() => { match message { - Some(message) => match message.event { + Some(Ok(message)) => match message.event { Some(boxlite_shared::exec_output::Event::Stdout(stdout)) => { - let _ = session_handle.data(channel_id, stdout.data).await; + if let Some(lost_bytes) = output_gap( + &mut stdout_offset, + stdout.offset, + stdout.data.len(), + stdout.total_bytes, + ) { + match output_gap_disposition(completion) { + OutputGapDisposition::Report => { + let _ = session_handle + .data( + channel_id, + output_gap_message("stdout", lost_bytes), + ) + .await; + } + OutputGapDisposition::Abort => { + warn!(%execution_id, lost_bytes, "SSH forwarding output gap"); + terminate_process_group(server.clone(), execution_id.clone()).await; + finish_channel( + &session_handle, + channel_id, + completion, + ExitNotification::Status(INDETERMINATE_EXIT_STATUS), + ) + .await; + return; + } + } + } + if !stdout.data.is_empty() { + let _ = session_handle.data(channel_id, stdout.data).await; + } } Some(boxlite_shared::exec_output::Event::Stderr(stderr)) => { - let _ = session_handle.extended_data(channel_id, 1, stderr.data).await; + if let Some(lost_bytes) = output_gap( + &mut stderr_offset, + stderr.offset, + stderr.data.len(), + stderr.total_bytes, + ) { + match output_gap_disposition(completion) { + OutputGapDisposition::Report => { + let _ = session_handle + .extended_data( + channel_id, + 1, + output_gap_message("stderr", lost_bytes), + ) + .await; + } + OutputGapDisposition::Abort => { + warn!(%execution_id, lost_bytes, "SSH forwarding output gap"); + terminate_process_group(server.clone(), execution_id.clone()).await; + finish_channel( + &session_handle, + channel_id, + completion, + ExitNotification::Status(INDETERMINATE_EXIT_STATUS), + ) + .await; + return; + } + } + } + if !stderr.data.is_empty() { + let _ = session_handle.extended_data(channel_id, 1, stderr.data).await; + } } None => {} }, + Some(Err(error)) => { + warn!(%error, %execution_id, "SSH output stream failed"); + terminate_process_group(server.clone(), execution_id.clone()).await; + finish_channel( + &session_handle, + channel_id, + completion, + ExitNotification::Status(INDETERMINATE_EXIT_STATUS), + ) + .await; + return; + } None => break, } } @@ -702,6 +797,59 @@ fn spawn_output_pump( }) } +fn output_gap( + expected_offset: &mut u64, + offset: Option, + data_len: usize, + total_bytes: Option, +) -> Option { + if let Some(total_bytes) = total_bytes { + if offset != Some(total_bytes) || total_bytes < *expected_offset { + warn!( + ?offset, + total_bytes, expected_offset, "invalid SSH execution output end frame" + ); + return None; + } + let lost_bytes = total_bytes - *expected_offset; + *expected_offset = total_bytes; + return (lost_bytes > 0).then_some(lost_bytes); + } + + let Some(offset) = offset else { + *expected_offset += data_len as u64; + return None; + }; + if offset < *expected_offset { + warn!( + offset, + expected_offset, "SSH execution output chunk moved backwards" + ); + return None; + } + + let lost_bytes = offset - *expected_offset; + *expected_offset = offset + data_len as u64; + (lost_bytes > 0).then_some(lost_bytes) +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum OutputGapDisposition { + Report, + Abort, +} + +fn output_gap_disposition(completion: ChannelCompletion) -> OutputGapDisposition { + match completion { + ChannelCompletion::Session => OutputGapDisposition::Report, + ChannelCompletion::Forwarding => OutputGapDisposition::Abort, + } +} + +fn output_gap_message(source: &str, lost_bytes: u64) -> Vec { + format!("[boxlite] {source} output dropped {lost_bytes} bytes\r\n").into_bytes() +} + /// Keep one lifecycle owner for every SSH-created execution. /// /// The output task may finish early because Attach failed or the SSH channel @@ -885,7 +1033,8 @@ mod tests { let (_stdin_peer, stdin) = pipe().unwrap(); let (stdout, _stdout_peer) = pipe().unwrap(); let (stderr, _stderr_peer) = pipe().unwrap(); - let handle = ExecHandle::new(Pid::from_raw(pid), stdin, stdout, Some(stderr)); + let handle = ExecHandle::new(Pid::from_raw(pid), stdin, stdout, Some(stderr)) + .expect("test pipe must register with Tokio"); let (exit, exit_tx) = ExitSlot::pending_for_test(); let state = ExecutionState::new_for_test(handle, exit); registry.register(execution_id.into(), state).await; @@ -1045,14 +1194,29 @@ mod tests { assert!(!ChannelCompletion::Forwarding.sends_exit_notification()); } + #[test] + fn output_gap_detects_mid_stream_and_terminal_loss() { + let mut offset = 0; + assert_eq!(output_gap(&mut offset, Some(0), 3, None), None); + assert_eq!(offset, 3); + assert_eq!(output_gap(&mut offset, Some(5), 2, None), Some(2)); + assert_eq!(offset, 7); + assert_eq!(output_gap(&mut offset, Some(10), 0, Some(10)), Some(3)); + assert_eq!(offset, 10); + } + #[test] fn streamlocal_readiness_prefix_can_be_fragmented() { let magic = crate::service::ssh::streamlocal::STREAMLOCAL_READY_MAGIC; let mut parser = StreamlocalReadyParser::default(); - assert_eq!(parser.push_stdout(&magic[..3]).unwrap(), None); - assert_eq!(parser.push_stdout(&magic[3..11]).unwrap(), None); - assert_eq!(parser.push_stdout(&magic[11..]).unwrap(), Some(Vec::new())); + assert_eq!(parser.push_stdout(Some(0), &magic[..3]).unwrap(), None); + assert_eq!(parser.push_stdout(Some(3), &magic[3..11]).unwrap(), None); + assert_eq!( + parser.push_stdout(Some(11), &magic[11..]).unwrap(), + Some(Vec::new()) + ); + assert_eq!(parser.stdout_offset, magic.len() as u64); } #[test] @@ -1063,15 +1227,49 @@ mod tests { let mut parser = StreamlocalReadyParser::default(); assert_eq!( - parser.push_stdout(&frame).unwrap(), + parser.push_stdout(Some(0), &frame).unwrap(), Some(b"socket payload".to_vec()) ); + assert_eq!(parser.stdout_offset, frame.len() as u64); } #[test] fn streamlocal_readiness_rejects_non_helper_output() { let mut parser = StreamlocalReadyParser::default(); - assert!(parser.push_stdout(b"not readiness").is_err()); + assert!(parser.push_stdout(Some(0), b"not readiness").is_err()); + } + + #[test] + fn streamlocal_ready_prefix_does_not_create_an_output_gap() { + let magic = crate::service::ssh::streamlocal::STREAMLOCAL_READY_MAGIC; + let mut parser = StreamlocalReadyParser::default(); + assert_eq!( + parser.push_stdout(Some(0), magic).unwrap(), + Some(Vec::new()) + ); + + let mut stdout_offset = parser.stdout_offset; + assert_eq!( + output_gap( + &mut stdout_offset, + Some(magic.len() as u64), + b"socket payload".len(), + None, + ), + None + ); + } + + #[test] + fn forwarding_output_gaps_abort_instead_of_writing_to_the_byte_stream() { + assert_eq!( + output_gap_disposition(ChannelCompletion::Session), + OutputGapDisposition::Report + ); + assert_eq!( + output_gap_disposition(ChannelCompletion::Forwarding), + OutputGapDisposition::Abort + ); } #[test] diff --git a/src/guest/src/service/ssh/reverse_streamlocal.rs b/src/guest/src/service/ssh/reverse_streamlocal.rs index 80412f783..5fe2e630e 100644 --- a/src/guest/src/service/ssh/reverse_streamlocal.rs +++ b/src/guest/src/service/ssh/reverse_streamlocal.rs @@ -513,7 +513,7 @@ struct RunningHelper { execution_id: String, stdin: mpsc::Sender, stdin_task: JoinHandle<()>, - output: mpsc::Receiver, + output: mpsc::Receiver>, buffered_stdout: Vec, } @@ -765,10 +765,10 @@ fn spawn_listener( } output = helper.output.recv() => { match output { - Some(ExecOutput { event: Some(exec_output::Event::Stderr(stderr)) }) => { + Some(Ok(ExecOutput { event: Some(exec_output::Event::Stderr(stderr)) })) => { debug!(bytes = stderr.data.len(), "reverse streamlocal helper stderr"); } - Some(ExecOutput { event: Some(exec_output::Event::Stdout(stdout)) }) => { + Some(Ok(ExecOutput { event: Some(exec_output::Event::Stdout(stdout)) })) => { if !append_stopped_marker_prefix( &mut helper.buffered_stdout, &stdout.data, @@ -777,7 +777,11 @@ fn spawn_listener( break End::HelperEnded; } } - Some(ExecOutput { event: None }) => {} + Some(Ok(ExecOutput { event: None })) => {} + Some(Err(error)) => { + warn!(%error, "reverse streamlocal helper output failed"); + break End::HelperEnded; + } None => break End::HelperEnded, } } @@ -923,7 +927,7 @@ impl MarkerParser { } async fn read_marker( - output: &mut mpsc::Receiver, + output: &mut mpsc::Receiver>, marker: &[u8], initial_stdout: Vec, max_trailing_bytes: usize, @@ -936,9 +940,13 @@ async fn read_marker( } } loop { - let message = output.recv().await.ok_or_else(|| { - "reverse streamlocal helper exited before control marker".to_string() - })?; + let message = output + .recv() + .await + .ok_or_else(|| { + "reverse streamlocal helper exited before control marker".to_string() + })? + .map_err(|error| format!("reverse streamlocal helper output failed: {error}"))?; match message.event { Some(exec_output::Event::Stdout(stdout)) => { if let Some(buffered) = parser.push(marker, &stdout.data, max_trailing_bytes)? { @@ -976,7 +984,7 @@ fn spawn_failed_helper_cleanup( server: Arc, registry: ExecutionRegistry, execution_id: String, - output: Option>, + output: Option>>, stdin_task: Option>, ) { spawn_execution_cleanup(server, registry, execution_id, output, stdin_task, true); @@ -986,7 +994,7 @@ fn spawn_execution_cleanup( server: Arc, registry: ExecutionRegistry, execution_id: String, - output: Option>, + output: Option>>, stdin_task: Option>, force_termination: bool, ) { @@ -994,6 +1002,9 @@ fn spawn_execution_cleanup( let output_task = output.map(|mut output| { tokio::spawn(async move { while let Some(message) = output.recv().await { + let Ok(message) = message else { + break; + }; if let Some(exec_output::Event::Stderr(stderr)) = message.event { debug!( bytes = stderr.data.len(), diff --git a/src/shared/proto/boxlite/v1/service.proto b/src/shared/proto/boxlite/v1/service.proto index 39c682ea2..5f31c0037 100644 --- a/src/shared/proto/boxlite/v1/service.proto +++ b/src/shared/proto/boxlite/v1/service.proto @@ -444,6 +444,7 @@ message AttachRequest { } message ExecOutput { + // An empty stdout/stderr event with total_bytes marks that stream's end. oneof event { Stdout stdout = 1; Stderr stderr = 2; @@ -452,10 +453,14 @@ message ExecOutput { message Stdout { bytes data = 1; + optional uint64 offset = 2; + optional uint64 total_bytes = 3; } message Stderr { bytes data = 1; + optional uint64 offset = 2; + optional uint64 total_bytes = 3; } // SendInput: client streaming stdin