diff --git a/Cargo.lock b/Cargo.lock index 2cb71b1e1..079c90333 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3810,6 +3810,7 @@ name = "microsandbox-agentd" version = "0.6.9" dependencies = [ "base64 0.23.1", + "bytes", "chrono", "ciborium", "libc", @@ -4048,6 +4049,7 @@ dependencies = [ name = "microsandbox-protocol" version = "0.6.9" dependencies = [ + "bytes", "chrono", "ciborium", "microsandbox-types", @@ -4244,8 +4246,7 @@ dependencies = [ [[package]] name = "msb_krun" version = "0.1.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "40e6b83f6cf90f676779409ac026f4429adfd7a11886085a59773ddd3aa8207f" +source = "git+https://github.com/superradcompany/libkrun?rev=ff087a1a0add4006f7ad4d753316cde1015a28b2#ff087a1a0add4006f7ad4d753316cde1015a28b2" dependencies = [ "crossbeam-channel", "kvm-bindings", @@ -4264,8 +4265,7 @@ dependencies = [ [[package]] name = "msb_krun_arch" version = "0.1.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "929f9765edd9e12cc45e5876998456bf2dfb0d2244d4fe6ddf2fdb5c3dd7e99c" +source = "git+https://github.com/superradcompany/libkrun?rev=ff087a1a0add4006f7ad4d753316cde1015a28b2#ff087a1a0add4006f7ad4d753316cde1015a28b2" dependencies = [ "kvm-bindings", "kvm-ioctls", @@ -4279,14 +4279,12 @@ dependencies = [ [[package]] name = "msb_krun_arch_gen" version = "0.1.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df83bf780774e4e131547a4e17d0ca389e5bbb04cdebcb626880b090cae20f2c" +source = "git+https://github.com/superradcompany/libkrun?rev=ff087a1a0add4006f7ad4d753316cde1015a28b2#ff087a1a0add4006f7ad4d753316cde1015a28b2" [[package]] name = "msb_krun_cpuid" version = "0.1.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c569d5d7285ed8f8f80e58454fb3ee213fd2c06ec647cc108a16a2e095558bb" +source = "git+https://github.com/superradcompany/libkrun?rev=ff087a1a0add4006f7ad4d753316cde1015a28b2#ff087a1a0add4006f7ad4d753316cde1015a28b2" dependencies = [ "kvm-bindings", "kvm-ioctls", @@ -4296,8 +4294,7 @@ dependencies = [ [[package]] name = "msb_krun_devices" version = "0.1.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "422da6e72923c2debae8eaeddd2990b924df9042079b51a323b4f3de37ebbdaf" +source = "git+https://github.com/superradcompany/libkrun?rev=ff087a1a0add4006f7ad4d753316cde1015a28b2#ff087a1a0add4006f7ad4d753316cde1015a28b2" dependencies = [ "bitflags 1.3.2", "capng", @@ -4326,8 +4323,7 @@ dependencies = [ [[package]] name = "msb_krun_hvf" version = "0.1.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5034e7b5ebd4f2f8caefb199f5f18483f8ed0c5cc6afb944b6ed0d127d59aa9" +source = "git+https://github.com/superradcompany/libkrun?rev=ff087a1a0add4006f7ad4d753316cde1015a28b2#ff087a1a0add4006f7ad4d753316cde1015a28b2" dependencies = [ "crossbeam-channel", "libloading 0.8.9", @@ -4338,8 +4334,7 @@ dependencies = [ [[package]] name = "msb_krun_kernel" version = "0.1.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04163d07028500bf08e6247c40c64dd5bae462bbb1997cf6a9613d019976a055" +source = "git+https://github.com/superradcompany/libkrun?rev=ff087a1a0add4006f7ad4d753316cde1015a28b2#ff087a1a0add4006f7ad4d753316cde1015a28b2" dependencies = [ "msb-vm-memory", "msb_krun_utils", @@ -4348,8 +4343,7 @@ dependencies = [ [[package]] name = "msb_krun_polly" version = "0.1.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9812193639b368ab800b5f87e18509ad30fd85a2e6a77a05265ffadfed2553d1" +source = "git+https://github.com/superradcompany/libkrun?rev=ff087a1a0add4006f7ad4d753316cde1015a28b2#ff087a1a0add4006f7ad4d753316cde1015a28b2" dependencies = [ "libc", "msb_krun_utils", @@ -4358,8 +4352,7 @@ dependencies = [ [[package]] name = "msb_krun_smbios" version = "0.1.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cfd3f05f6447f0ad85f04f688b106eb9f04bb64525fee1183757cf9705070a0" +source = "git+https://github.com/superradcompany/libkrun?rev=ff087a1a0add4006f7ad4d753316cde1015a28b2#ff087a1a0add4006f7ad4d753316cde1015a28b2" dependencies = [ "msb-vm-memory", ] @@ -4367,8 +4360,7 @@ dependencies = [ [[package]] name = "msb_krun_utils" version = "0.1.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b5ce62484f07792ab43a5ac64baa84435556f56a7633805e15030e990c3b98a3" +source = "git+https://github.com/superradcompany/libkrun?rev=ff087a1a0add4006f7ad4d753316cde1015a28b2#ff087a1a0add4006f7ad4d753316cde1015a28b2" dependencies = [ "bitflags 1.3.2", "crossbeam-channel", @@ -4383,8 +4375,7 @@ dependencies = [ [[package]] name = "msb_krun_vmm" version = "0.1.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67546eb70fa306e3c7e97f525dac2db44b9cc1642b6684aec4369e3b7f9b3af3" +source = "git+https://github.com/superradcompany/libkrun?rev=ff087a1a0add4006f7ad4d753316cde1015a28b2#ff087a1a0add4006f7ad4d753316cde1015a28b2" dependencies = [ "bzip2", "crossbeam-channel", diff --git a/Cargo.toml b/Cargo.toml index e061ab3ab..0bdf85de9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -88,8 +88,9 @@ microsandbox-protocol = { version = "=0.6.9", path = "crates/protocol" } microsandbox-runtime = { version = "=0.6.9", path = "crates/runtime", default-features = false } microsandbox-utils = { version = "=0.6.9", path = "crates/utils" } microsandbox-vsock = { version = "=0.6.9", path = "crates/vsock" } -msb_krun = "=0.1.31" -msb_krun_utils = "=0.1.31" +# Keep CI on the exact prerequisite revision until its console API is released to crates.io. +msb_krun = { git = "https://github.com/superradcompany/libkrun", rev = "ff087a1a0add4006f7ad4d753316cde1015a28b2" } +msb_krun_utils = { git = "https://github.com/superradcompany/libkrun", rev = "ff087a1a0add4006f7ad4d753316cde1015a28b2" } test-macros = { path = "crates/testing/macros" } test-utils = { path = "crates/testing/utils" } diff --git a/crates/agentd/Cargo.toml b/crates/agentd/Cargo.toml index 52ffe6d34..35b7dfbc6 100644 --- a/crates/agentd/Cargo.toml +++ b/crates/agentd/Cargo.toml @@ -16,6 +16,7 @@ path = "lib/lib.rs" [dependencies] base64.workspace = true +bytes.workspace = true chrono.workspace = true ciborium.workspace = true libc.workspace = true diff --git a/crates/agentd/lib/agent.rs b/crates/agentd/lib/agent.rs index 6358f3a68..64ef40f8e 100644 --- a/crates/agentd/lib/agent.rs +++ b/crates/agentd/lib/agent.rs @@ -8,9 +8,10 @@ use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Instant; +use bytes::BytesMut; use chrono::Utc; use tokio::io::unix::AsyncFd; -use tokio::sync::{mpsc, watch}; +use tokio::sync::watch; use tokio::time::{self, Duration}; use microsandbox_protocol::HANDOFF_POWEROFF_TIMEOUT; @@ -34,7 +35,8 @@ use crate::fs::{FsReadSession, FsState, FsStreamSession, FsWriteSession}; use crate::process::ProcessManager; use crate::serial::AGENT_PORT_NAME; use crate::session::{ - ExecSession, RawActivity, RawSessionCompletion, SessionOutput, resolve_default_user, + ExecSession, RawActivity, RawSessionCompletion, SessionOutput, SessionOutputSender, + resolve_default_user, }; use crate::tcp::TcpSession; use crate::{clock, fs, handoff, heartbeat, serial}; @@ -155,13 +157,13 @@ pub async fn run( // Buffer for serial reads. let mut read_buf = vec![0u8; SERIAL_READ_BUF_SIZE]; - let mut serial_in_buf = Vec::new(); + let mut serial_in_buf = BytesMut::new(); let mut serial_out_buf = Vec::new(); let mut state = AgentState::default(); // Channel for session output events. - let (session_tx, mut session_rx) = mpsc::unbounded_channel::<(u32, SessionOutput)>(); + let (session_tx, mut session_rx) = SessionOutputSender::channel(); // Heartbeat/activity state. let mut activity = ActivityTracker::new(); @@ -238,19 +240,9 @@ pub async fn run( // message-level failures are reported on the same // correlation ID with `core.error`; unrecoverable // frame-level failures still close the agent loop. - while let Some(frame) = codec::try_decode_raw_from_buf(&mut serial_in_buf) + while let Some(msg) = codec::try_decode_from_bytes(&mut serial_in_buf) .map_err(|e| AgentdError::ExecSession(format!("decode frame: {e}")))? { - let id = frame.id; - let msg = match codec::raw_frame_to_message(frame) { - Ok(msg) => msg, - Err(e) => { - return Err(AgentdError::ExecSession(format!( - "decode message for id {id}: {e}" - ))); - } - }; - if msg.flags != msg.t.flags() { let out_before = serial_out_buf.len(); encode_core_error_if_supported( @@ -316,8 +308,9 @@ pub async fn run( } // Receive output events from session reader tasks. - Some((id, output)) = session_rx.recv() => { - match output { + Some(envelope) = session_rx.recv() => { + let id = envelope.id; + match envelope.output { SessionOutput::Stdout(data) => { let len = data.len(); let msg = Message::with_payload(MessageType::ExecStdout, id, &ExecStdout { data }) @@ -352,8 +345,13 @@ pub async fn run( &mut state.read_sessions, &mut state.tcp_sessions, ); - // Pre-encoded frame — write directly to output buffer. - serial_out_buf.extend_from_slice(&output.frame); + // The producer already owns an encoded frame. Write from that allocation + // directly so multi-megabyte FS/TCP frames are not copied into a second + // serial staging buffer. + if !serial_out_buf.is_empty() { + flush_write_buf(&async_port, &mut serial_out_buf).await?; + } + write_all_async_fd(&async_port, &output.frame).await?; } } publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity); @@ -413,7 +411,7 @@ async fn handle_message( msg: Message, state: &mut AgentState, activity: &mut ActivityTracker, - session_tx: &mpsc::UnboundedSender<(u32, SessionOutput)>, + session_tx: &SessionOutputSender, out_buf: &mut Vec, config: &AgentdConfig, ) -> AgentdResult<()> { @@ -1171,11 +1169,22 @@ fn write_all_to_fd(fd: i32, mut buf: &[u8], deadline: Instant) -> AgentdResult<( /// Flushes the write buffer to the async fd. async fn flush_write_buf(fd: &AsyncFd, buf: &mut Vec) -> AgentdResult<()> { - while !buf.is_empty() { + write_all_async_fd(fd, buf).await?; + buf.clear(); + Ok(()) +} + +/// Write an immutable region to the nonblocking serial descriptor with cursor advancement. +async fn write_all_async_fd(fd: &AsyncFd, buf: &[u8]) -> AgentdResult<()> { + let mut written = 0; + while written < buf.len() { let mut guard = fd.writable().await?; - match guard.try_io(|inner| write_to_fd(inner.get_ref().as_raw_fd(), buf)) { + match guard.try_io(|inner| write_to_fd(inner.get_ref().as_raw_fd(), &buf[written..])) { Ok(Ok(n)) => { - buf.drain(..n); + if n == 0 { + return Err(std::io::Error::from(std::io::ErrorKind::WriteZero).into()); + } + written += n; } Ok(Err(e)) if e.kind() == std::io::ErrorKind::Interrupted => continue, Ok(Err(e)) => return Err(e.into()), diff --git a/crates/agentd/lib/fs.rs b/crates/agentd/lib/fs.rs index f449d2350..37f6be190 100644 --- a/crates/agentd/lib/fs.rs +++ b/crates/agentd/lib/fs.rs @@ -19,10 +19,12 @@ use microsandbox_protocol::fs::{ }; use microsandbox_protocol::message::{Message, MessageType}; use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; -use tokio::sync::{Mutex, mpsc}; +use tokio::sync::Mutex; use tokio::task::JoinHandle; -use crate::session::{RawActivity, RawSessionCompletion, RawSessionOutput, SessionOutput}; +use crate::session::{ + RawActivity, RawSessionCompletion, RawSessionOutput, SessionOutput, SessionOutputSender, +}; //-------------------------------------------------------------------------------------------------- // Constants @@ -302,7 +304,7 @@ pub async fn handle_fs_request( req: FsRequest, state: &mut FsState, out_buf: &mut Vec, - session_tx: &mpsc::UnboundedSender<(u32, SessionOutput)>, + session_tx: &SessionOutputSender, ) -> Result, String> { match req.op { FsOp::RealPath { path } => { @@ -749,11 +751,11 @@ async fn handle_read_stream( file: Arc>, offset: u64, len: Option, - tx: &mpsc::UnboundedSender<(u32, SessionOutput)>, + tx: &SessionOutputSender, ) { let mut file = file.lock().await; if let Err(e) = file.seek(std::io::SeekFrom::Start(offset)).await { - send_raw_response(id, false, Some(format!("seek: {e}")), None, tx); + send_raw_response(id, false, Some(format!("seek: {e}")), None, tx).await; return; } @@ -762,6 +764,11 @@ async fn handle_read_stream( let mut buf = Vec::new(); loop { + // Reserve before the file read and CBOR materialization so concurrent streams cannot + // each create an uncharged maximum-sized frame while aggregate output is saturated. + let Some(permit) = tx.reserve(codec::MAX_FRAME_SIZE as usize + 4).await else { + return; + }; let read_len = match remaining { Some(0) => break, Some(n) => chunk.len().min(n as usize), @@ -780,7 +787,8 @@ async fn handle_read_stream( let msg = match Message::with_payload(MessageType::FsData, id, &data) { Ok(msg) => msg, Err(e) => { - send_raw_response(id, false, Some(format!("encode chunk: {e}")), None, tx); + send_raw_response(id, false, Some(format!("encode chunk: {e}")), None, tx) + .await; return; } }; @@ -792,22 +800,27 @@ async fn handle_read_stream( Some(format!("encode chunk frame: {e}")), None, tx, - ); + ) + .await; return; } - let output = RawSessionOutput::new(buf.clone(), RawActivity::fs_bytes(n), None); - if tx.send((id, SessionOutput::Raw(output))).is_err() { + let output = + RawSessionOutput::new(std::mem::take(&mut buf), RawActivity::fs_bytes(n), None); + if !tx + .send_reserved(id, SessionOutput::Raw(output), permit) + .await + { return; } } Err(e) => { - send_raw_response(id, false, Some(format!("read: {e}")), None, tx); + send_raw_response(id, false, Some(format!("read: {e}")), None, tx).await; return; } } } - send_raw_response(id, true, None, None, tx); + send_raw_response(id, true, None, None, tx).await; } //-------------------------------------------------------------------------------------------------- @@ -978,12 +991,12 @@ fn encode_response(id: u32, resp: FsResponse, out_buf: &mut Vec) -> Result<( Ok(()) } -fn send_raw_response( +async fn send_raw_response( id: u32, ok: bool, error: Option, data: Option, - tx: &mpsc::UnboundedSender<(u32, SessionOutput)>, + tx: &SessionOutputSender, ) { let resp = FsResponse { ok, error, data }; match Message::with_payload(MessageType::FsResponse, id, &resp) { @@ -996,7 +1009,7 @@ fn send_raw_response( RawActivity::guest_message(), Some(RawSessionCompletion::FsRead), ); - let _ = tx.send((id, SessionOutput::Raw(output))); + let _ = tx.send(id, SessionOutput::Raw(output)).await; } Err(e) => { eprintln!("failed to encode fs response frame for {id}: {e}"); diff --git a/crates/agentd/lib/session.rs b/crates/agentd/lib/session.rs index 869b70185..d9a697c1e 100644 --- a/crates/agentd/lib/session.rs +++ b/crates/agentd/lib/session.rs @@ -11,7 +11,7 @@ use std::{iter, mem, ptr}; use nix::pty; use nix::sys::signal::Signal; use tokio::io::AsyncReadExt; -use tokio::sync::mpsc; +use tokio::sync::{Semaphore, mpsc}; use microsandbox_protocol::exec::{ExecFailed, ExecFailureKind, ExecRequest}; @@ -32,6 +32,15 @@ const PR_CAP_AMBIENT: libc::c_int = 47; const PR_CAP_AMBIENT_CLEAR_ALL: libc::c_int = 4; const DEFAULT_USER_SPEC: &str = "0:0"; +/// Aggregate guest-to-host data retained outside the serial output buffer. +const SESSION_OUTPUT_BYTE_CAPACITY: usize = 32 * 1024 * 1024; + +/// Allocation granularity used by the data budget. +const SESSION_OUTPUT_BUDGET_GRANULE: usize = 4096; + +/// Maximum number of data or control events waiting for the serial writer. +const SESSION_OUTPUT_ITEM_CAPACITY: usize = 1024; + //-------------------------------------------------------------------------------------------------- // Functions: classify //-------------------------------------------------------------------------------------------------- @@ -148,6 +157,28 @@ pub enum SessionOutput { Raw(RawSessionOutput), } +/// One queued session event and the data-budget capacity owned by its buffer. +pub struct SessionOutputEnvelope { + /// Correlation ID for the session event. + pub id: u32, + + /// Event consumed by the main serial loop. + pub output: SessionOutput, + + /// Capacity follows the allocation and is released only after serial output consumes it. + _permit: Option, +} + +/// Capacity reserved before a producer reads or encodes a data-bearing event. +pub struct SessionOutputPermit(tokio::sync::OwnedSemaphorePermit); + +/// Cloneable producer for the byte-bounded session output queue. +#[derive(Clone)] +pub struct SessionOutputSender { + tx: mpsc::Sender, + data_budget: Arc, +} + /// Pre-encoded session output plus the accounting metadata known by its producer. pub struct RawSessionOutput { /// Encoded protocol frame bytes. @@ -248,6 +279,93 @@ impl RawSessionOutput { } } +impl SessionOutput { + /// Bytes retained by this event that count against bulk output capacity. + fn budget_bytes(&self) -> usize { + let allocation = match self { + Self::Stdout(data) | Self::Stderr(data) => data.capacity(), + Self::Raw(output) + if output.activity.fs_bytes != 0 || output.activity.tcp_bytes != 0 => + { + output.frame.capacity() + } + Self::Exited(_) | Self::Raw(_) => 0, + }; + + allocation + .checked_add(SESSION_OUTPUT_BUDGET_GRANULE - 1) + .map(|bytes| bytes / SESSION_OUTPUT_BUDGET_GRANULE * SESSION_OUTPUT_BUDGET_GRANULE) + .unwrap_or(usize::MAX) + } +} + +impl SessionOutputSender { + /// Create one ordered queue with a separate byte budget for data-bearing events. + pub fn channel() -> (Self, mpsc::Receiver) { + let (tx, rx) = mpsc::channel(SESSION_OUTPUT_ITEM_CAPACITY); + ( + Self { + tx, + data_budget: Arc::new(Semaphore::new(SESSION_OUTPUT_BYTE_CAPACITY)), + }, + rx, + ) + } + + /// Queue an event after its retained allocation has acquired aggregate capacity. + pub async fn send(&self, id: u32, output: SessionOutput) -> bool { + let budget_bytes = output.budget_bytes(); + let Some(permit) = self.reserve(budget_bytes).await else { + eprintln!("agentd session output {id} exceeds byte budget: {budget_bytes} bytes"); + return false; + }; + + self.send_reserved(id, output, permit).await + } + + /// Reserve capacity before reading or encoding up to `max_bytes` of output. + pub async fn reserve(&self, max_bytes: usize) -> Option { + let budget_bytes = max_bytes.checked_add(SESSION_OUTPUT_BUDGET_GRANULE - 1)? + / SESSION_OUTPUT_BUDGET_GRANULE + * SESSION_OUTPUT_BUDGET_GRANULE; + if budget_bytes > SESSION_OUTPUT_BYTE_CAPACITY { + return None; + } + let permit_count = u32::try_from(budget_bytes).ok()?; + Arc::clone(&self.data_budget) + .acquire_many_owned(permit_count) + .await + .ok() + .map(SessionOutputPermit) + } + + /// Queue output using capacity obtained before the producer created its allocation. + pub async fn send_reserved( + &self, + id: u32, + output: SessionOutput, + permit: SessionOutputPermit, + ) -> bool { + let charged = output.budget_bytes(); + if charged > permit.0.num_permits() { + eprintln!( + "agentd session output {id} exceeded its reservation: {charged} > {} bytes", + permit.0.num_permits() + ); + return false; + } + + self.tx + .send(SessionOutputEnvelope { + id, + output, + _permit: (charged != 0).then_some(permit.0), + }) + .await + .is_ok() + } +} + impl RawActivity { /// A guest-to-host frame with no byte counter. pub fn guest_message() -> Self { @@ -284,7 +402,7 @@ impl ExecSession { pub fn spawn( id: u32, req: &ExecRequest, - tx: mpsc::UnboundedSender<(u32, SessionOutput)>, + tx: SessionOutputSender, default_user: Option<&str>, security_profile: SecurityProfile, ) -> AgentdResult { @@ -371,7 +489,7 @@ impl ExecSession { fn spawn_pty( id: u32, req: &ExecRequest, - tx: mpsc::UnboundedSender<(u32, SessionOutput)>, + tx: SessionOutputSender, default_user: Option<&str>, security_profile: SecurityProfile, process_manager: &Arc, @@ -583,7 +701,7 @@ impl ExecSession { fn spawn_pipe( id: u32, req: &ExecRequest, - tx: mpsc::UnboundedSender<(u32, SessionOutput)>, + tx: SessionOutputSender, default_user: Option<&str>, security_profile: SecurityProfile, process_manager: &Arc, @@ -1165,9 +1283,10 @@ async fn pty_reader_task( id: u32, master_fd: OwnedFd, exit_watcher: ProcessExitWatcher, - tx: mpsc::UnboundedSender<(u32, SessionOutput)>, + tx: SessionOutputSender, ) { let tx_output = tx.clone(); + let runtime_handle = tokio::runtime::Handle::current(); let read_result = tokio::task::spawn_blocking(move || { // PTY masters are safer with a dedicated blocking read loop than with // edge-driven readiness. Fast writers followed by process exit can @@ -1183,10 +1302,16 @@ async fn pty_reader_task( let n = unsafe { libc::read(raw, buf.as_mut_ptr() as *mut libc::c_void, buf.len()) }; if n > 0 { - if tx_output - .send((id, SessionOutput::Stdout(buf[..n as usize].to_vec()))) - .is_err() - { + let n = n as usize; + let sent = runtime_handle.block_on(async { + let Some(permit) = tx_output.reserve(n).await else { + return false; + }; + tx_output + .send_reserved(id, SessionOutput::Stdout(buf[..n].to_vec()), permit) + .await + }); + if !sent { break; } continue; @@ -1209,7 +1334,7 @@ async fn pty_reader_task( let _ = read_result; let code = exit_watcher.await; - let _ = tx.send((id, SessionOutput::Exited(code))); + let _ = tx.send(id, SessionOutput::Exited(code)).await; } /// Background task that reads from piped stdout/stderr and sends output events. @@ -1218,7 +1343,7 @@ async fn pipe_reader_task( stdout: Option, stderr: Option, exit_watcher: ProcessExitWatcher, - tx: mpsc::UnboundedSender<(u32, SessionOutput)>, + tx: SessionOutputSender, ) { let mut stdout = stdout; let mut stderr = stderr; @@ -1242,7 +1367,19 @@ async fn pipe_reader_task( stdout_eof = true; } Ok(n) => { - let _ = tx.send((id, SessionOutput::Stdout(stdout_buf[..n].to_vec()))); + let Some(permit) = tx.reserve(n).await else { + break; + }; + if !tx + .send_reserved( + id, + SessionOutput::Stdout(stdout_buf[..n].to_vec()), + permit, + ) + .await + { + break; + } } } } @@ -1258,7 +1395,19 @@ async fn pipe_reader_task( stderr_eof = true; } Ok(n) => { - let _ = tx.send((id, SessionOutput::Stderr(stderr_buf[..n].to_vec()))); + let Some(permit) = tx.reserve(n).await else { + break; + }; + if !tx + .send_reserved( + id, + SessionOutput::Stderr(stderr_buf[..n].to_vec()), + permit, + ) + .await + { + break; + } } } } @@ -1267,7 +1416,7 @@ async fn pipe_reader_task( let code = exit_watcher.await; - let _ = tx.send((id, SessionOutput::Exited(code))); + let _ = tx.send(id, SessionOutput::Exited(code)).await; } //-------------------------------------------------------------------------------------------------- @@ -1298,6 +1447,38 @@ mod tests { const RUNTIME_TEST_NAME: &str = "session::tests::test_spawn_survives_runtime_replacement"; const PIPE_OWNER_HELPER_ENV: &str = "MSB_AGENTD_PIPE_OWNER_HELPER"; const PIPE_OWNER_HELPER_SENTINEL: &str = "pipe-owner-helper-passed"; + + #[tokio::test] + async fn session_output_permit_lives_until_envelope_is_consumed() { + let (tx, mut rx) = SessionOutputSender::channel(); + assert!(tx.send(7, SessionOutput::Stdout(vec![0; 4096])).await); + assert_eq!( + tx.data_budget.available_permits(), + SESSION_OUTPUT_BYTE_CAPACITY - 4096 + ); + + let envelope = rx.recv().await.unwrap(); + assert_eq!( + tx.data_budget.available_permits(), + SESSION_OUTPUT_BYTE_CAPACITY - 4096 + ); + drop(envelope); + assert_eq!( + tx.data_budget.available_permits(), + SESSION_OUTPUT_BYTE_CAPACITY + ); + } + + #[tokio::test] + async fn control_output_remains_admissible_when_data_budget_is_exhausted() { + let (tx, mut rx) = SessionOutputSender::channel(); + let full_budget = tx.reserve(SESSION_OUTPUT_BYTE_CAPACITY).await.unwrap(); + + assert!(tx.send(9, SessionOutput::Exited(0)).await); + let envelope = rx.recv().await.unwrap(); + assert!(matches!(envelope.output, SessionOutput::Exited(0))); + drop(full_budget); + } const PIPE_OWNER_TEST_NAME: &str = "session::tests::test_piped_process_exit_outlives_spawning_runtime"; @@ -1347,7 +1528,7 @@ mod tests { std::io::Error::last_os_error() ); - let (tx, mut rx) = mpsc::unbounded_channel(); + let (tx, mut rx) = SessionOutputSender::channel(); let req = ExecRequest { cmd: "/bin/sh".to_string(), args: vec!["-c".to_string(), "sleep 30 & echo $!".to_string()], @@ -1366,9 +1547,9 @@ mod tests { let mut stdout = Vec::new(); time::timeout(Duration::from_secs(10), async { while !stdout.contains(&b'\n') { - let (id, output) = rx.recv().await.expect("session output"); - assert_eq!(id, 17); - match output { + let envelope = rx.recv().await.expect("session output"); + assert_eq!(envelope.id, 17); + match envelope.output { SessionOutput::Stdout(data) => stdout.extend_from_slice(&data), SessionOutput::Exited(code) => panic!("session exited early with {code}"), SessionOutput::Stderr(_) | SessionOutput::Raw(_) => {} @@ -1415,9 +1596,9 @@ mod tests { .expect("signal descendants through completed process registration"); let exit = time::timeout(Duration::from_secs(5), async { loop { - let (id, output) = rx.recv().await.expect("session output after signal"); - assert_eq!(id, 17); - if let SessionOutput::Exited(code) = output { + let envelope = rx.recv().await.expect("session output after signal"); + assert_eq!(envelope.id, 17); + if let SessionOutput::Exited(code) = envelope.output { break code; } } @@ -1484,7 +1665,7 @@ mod tests { const PROCESS_COUNT: u32 = 12; let runtime_handle = tokio::runtime::Handle::current(); - let (tx, mut rx) = mpsc::unbounded_channel(); + let (tx, mut rx) = SessionOutputSender::channel(); let mut spawn_threads = Vec::new(); for offset in 0..PROCESS_COUNT { let handle = runtime_handle.clone(); @@ -1521,9 +1702,9 @@ mod tests { let mut exits = HashMap::new(); time::timeout(Duration::from_secs(15), async { while exits.len() < PROCESS_COUNT as usize { - let (id, output) = rx.recv().await.expect("session output"); - if let SessionOutput::Exited(code) = output { - exits.insert(id, code); + let envelope = rx.recv().await.expect("session output"); + if let SessionOutput::Exited(code) = envelope.output { + exits.insert(envelope.id, code); } } }) @@ -1580,7 +1761,7 @@ mod tests { } async fn run_single_pipe_spawn(id: u32, code: i32) { - let (tx, mut rx) = mpsc::unbounded_channel(); + let (tx, mut rx) = SessionOutputSender::channel(); let req = ExecRequest { cmd: "/bin/sh".to_string(), args: vec!["-c".to_string(), format!("exit {code}")], @@ -1597,9 +1778,9 @@ mod tests { let actual = time::timeout(Duration::from_secs(5), async { loop { - let (actual_id, output) = rx.recv().await.expect("session output"); - assert_eq!(actual_id, id); - if let SessionOutput::Exited(actual) = output { + let envelope = rx.recv().await.expect("session output"); + assert_eq!(envelope.id, id); + if let SessionOutput::Exited(actual) = envelope.output { break actual; } } @@ -1677,7 +1858,7 @@ mod tests { #[tokio::test] async fn test_pty_reader_drains_ready_fd() { - let (tx, mut rx) = mpsc::unbounded_channel(); + let (tx, mut rx) = SessionOutputSender::channel(); let req = ExecRequest { cmd: "/bin/sh".to_string(), args: vec![ @@ -1700,9 +1881,9 @@ mod tests { let mut exit = None; let recv_result = time::timeout(Duration::from_secs(15), async { - while let Some((id, output)) = rx.recv().await { - assert_eq!(id, 7); - match output { + while let Some(envelope) = rx.recv().await { + assert_eq!(envelope.id, 7); + match envelope.output { SessionOutput::Stdout(data) => stdout.extend_from_slice(&data), SessionOutput::Exited(code) => { exit = Some(code); @@ -1891,7 +2072,7 @@ mod tests { #[tokio::test] async fn test_spawn_pipe_error_does_not_include_probe_details() { - let (tx, _rx) = mpsc::unbounded_channel(); + let (tx, _rx) = SessionOutputSender::channel(); let req = ExecRequest { cmd: "/definitely/not/a/real/binary".to_string(), args: Vec::new(), diff --git a/crates/agentd/lib/tcp.rs b/crates/agentd/lib/tcp.rs index 8036c1834..17ddcd400 100644 --- a/crates/agentd/lib/tcp.rs +++ b/crates/agentd/lib/tcp.rs @@ -14,7 +14,12 @@ use microsandbox_protocol::codec; use microsandbox_protocol::message::{Message, MessageType}; use microsandbox_protocol::tcp::{TcpClosed, TcpConnect, TcpConnected, TcpData, TcpEof, TcpFailed}; -use crate::session::{RawActivity, RawSessionCompletion, RawSessionOutput, SessionOutput}; +#[cfg(test)] +use crate::session::SessionOutputEnvelope; +use crate::session::{ + RawActivity, RawSessionCompletion, RawSessionOutput, SessionOutput, SessionOutputPermit, + SessionOutputSender, +}; //-------------------------------------------------------------------------------------------------- // Constants @@ -23,6 +28,12 @@ use crate::session::{RawActivity, RawSessionCompletion, RawSessionOutput, Sessio /// TCP stream read chunk size. const TCP_CHUNK_SIZE: usize = 64 * 1024; +/// Capacity reserved before cloning and encoding one TCP data chunk. +/// +/// The factor of two covers CBOR/framing overhead and allocator growth while keeping hundreds of +/// TCP chunks eligible under the aggregate 32 MiB output budget. +const TCP_OUTPUT_RESERVATION: usize = 2 * TCP_CHUNK_SIZE; + /// How many host->guest command frames may queue before the agent loop has to /// wait. Bounding this turns a slow or stalled destination into backpressure /// (the serial reader pauses, which throttles the SSH window) instead of @@ -105,11 +116,7 @@ impl TcpSession { /// `core.tcp.failed` on error/timeout over `session_tx`; the host correlates /// either reply by id. The returned session is live immediately, with /// commands queued until the connect completes. - pub fn open( - id: u32, - req: TcpConnect, - session_tx: &mpsc::UnboundedSender<(u32, SessionOutput)>, - ) -> Self { + pub fn open(id: u32, req: TcpConnect, session_tx: &SessionOutputSender) -> Self { let (commands_tx, commands_rx) = mpsc::channel(TCP_COMMAND_CAPACITY); let output_tx = session_tx.clone(); let task = tokio::spawn(async move { @@ -138,7 +145,7 @@ async fn connect_and_relay( id: u32, req: TcpConnect, commands: mpsc::Receiver, - tx: mpsc::UnboundedSender<(u32, SessionOutput)>, + tx: SessionOutputSender, ) { let connect = TcpStream::connect((req.host.as_str(), req.port)); let stream = match tokio::time::timeout(TCP_CONNECT_TIMEOUT, connect).await { @@ -153,7 +160,8 @@ async fn connect_and_relay( RawActivity::guest_message(), Some(RawSessionCompletion::Tcp), &tx, - ); + ) + .await; return; } Err(_elapsed) => { @@ -166,7 +174,8 @@ async fn connect_and_relay( RawActivity::guest_message(), Some(RawSessionCompletion::Tcp), &tx, - ); + ) + .await; return; } }; @@ -178,7 +187,9 @@ async fn connect_and_relay( RawActivity::guest_message(), None, &tx, - ) { + ) + .await + { return; } @@ -189,7 +200,7 @@ async fn relay_tcp_session( id: u32, mut stream: TcpStream, mut commands: mpsc::Receiver, - tx: mpsc::UnboundedSender<(u32, SessionOutput)>, + tx: SessionOutputSender, ) { let mut read_buf = vec![0u8; TCP_CHUNK_SIZE]; let mut terminal_sent = false; @@ -209,19 +220,17 @@ async fn relay_tcp_session( RawActivity::guest_message(), None, &tx, - ); + ) + .await; read_eof = true; } Ok(n) => { + let Some(permit) = tx.reserve(TCP_OUTPUT_RESERVATION).await else { + break; + }; let data = read_buf[..n].to_vec(); - if !send_raw_tcp_message( - id, - MessageType::TcpData, - &TcpData { data }, - RawActivity::tcp_bytes(n), - None, - &tx, - ) { + if !send_raw_tcp_data(id, data, n, permit, &tx).await + { break; } } @@ -235,7 +244,8 @@ async fn relay_tcp_session( RawActivity::guest_message(), Some(RawSessionCompletion::Tcp), &tx, - ); + ) + .await; break; } } @@ -253,7 +263,8 @@ async fn relay_tcp_session( RawActivity::guest_message(), Some(RawSessionCompletion::Tcp), &tx, - ); + ) + .await; break; } } @@ -268,7 +279,8 @@ async fn relay_tcp_session( RawActivity::guest_message(), Some(RawSessionCompletion::Tcp), &tx, - ); + ) + .await; break; } } @@ -288,7 +300,8 @@ async fn relay_tcp_session( RawActivity::guest_message(), Some(RawSessionCompletion::Tcp), &tx, - ); + ) + .await; } } @@ -303,22 +316,23 @@ fn encode_tcp_message( Ok(()) } -fn send_raw_tcp_message( +async fn send_raw_tcp_message( id: u32, t: MessageType, payload: &T, activity: RawActivity, completion: Option, - tx: &mpsc::UnboundedSender<(u32, SessionOutput)>, + tx: &SessionOutputSender, ) -> bool { let mut buf = Vec::new(); match encode_tcp_message(id, t, payload, &mut buf) { - Ok(()) => tx - .send(( + Ok(()) => { + tx.send( id, SessionOutput::Raw(RawSessionOutput::new(buf, activity, completion)), - )) - .is_ok(), + ) + .await + } Err(e) => { eprintln!("failed to encode tcp message for {id}: {e}"); false @@ -326,6 +340,35 @@ fn send_raw_tcp_message( } } +/// Encode a TCP data event only after its retained allocation has reserved capacity. +async fn send_raw_tcp_data( + id: u32, + data: Vec, + byte_count: usize, + permit: SessionOutputPermit, + tx: &SessionOutputSender, +) -> bool { + let mut buf = Vec::new(); + match encode_tcp_message(id, MessageType::TcpData, &TcpData { data }, &mut buf) { + Ok(()) => { + tx.send_reserved( + id, + SessionOutput::Raw(RawSessionOutput::new( + buf, + RawActivity::tcp_bytes(byte_count), + None, + )), + permit, + ) + .await + } + Err(error) => { + eprintln!("failed to encode TCP data for {id}: {error}"); + false + } + } +} + //-------------------------------------------------------------------------------------------------- // Tests //-------------------------------------------------------------------------------------------------- @@ -341,7 +384,7 @@ mod tests { #[tokio::test] async fn connect_failure_sends_terminal_failed() { - let (session_tx, mut session_rx) = mpsc::unbounded_channel(); + let (session_tx, mut session_rx) = SessionOutputSender::channel(); let session = TcpSession::open( 7, @@ -366,7 +409,7 @@ mod tests { async fn close_request_finishes_session_task() { let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); let port = listener.local_addr().unwrap().port(); - let (session_tx, mut session_rx) = mpsc::unbounded_channel(); + let (session_tx, mut session_rx) = SessionOutputSender::channel(); let accept_task = tokio::spawn(async move { let (_socket, _) = listener.accept().await.unwrap(); tokio::time::sleep(Duration::from_secs(5)).await; @@ -394,7 +437,7 @@ mod tests { async fn destination_eof_keeps_session_open_for_host_writes() { let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap(); let port = listener.local_addr().unwrap().port(); - let (session_tx, mut session_rx) = mpsc::unbounded_channel(); + let (session_tx, mut session_rx) = SessionOutputSender::channel(); // The destination half-closes its write side, then keeps reading so it // still receives whatever the host sends after the EOF. @@ -456,9 +499,9 @@ mod tests { codec::try_decode_from_buf(buf).unwrap().unwrap() } - async fn recv_message(rx: &mut mpsc::UnboundedReceiver<(u32, SessionOutput)>) -> Message { - let (_id, output) = rx.recv().await.unwrap(); - let SessionOutput::Raw(mut output) = output else { + async fn recv_message(rx: &mut mpsc::Receiver) -> Message { + let envelope = rx.recv().await.unwrap(); + let SessionOutput::Raw(mut output) = envelope.output else { panic!("expected SessionOutput::Raw frame"); }; decode_one_message(&mut output.frame) diff --git a/crates/protocol/Cargo.toml b/crates/protocol/Cargo.toml index d72dab01a..eef989479 100644 --- a/crates/protocol/Cargo.toml +++ b/crates/protocol/Cargo.toml @@ -11,6 +11,7 @@ edition.workspace = true path = "lib/lib.rs" [dependencies] +bytes.workspace = true chrono.workspace = true ciborium.workspace = true microsandbox-types.workspace = true diff --git a/crates/protocol/lib/codec.rs b/crates/protocol/lib/codec.rs index f1e3b8553..3c9eb9546 100644 --- a/crates/protocol/lib/codec.rs +++ b/crates/protocol/lib/codec.rs @@ -5,6 +5,9 @@ //! The correlation ID and flags sit in a fixed-position binary header so that //! relay intermediaries can route frames without CBOR parsing. +use std::io::IoSlice; + +use bytes::{Buf, BytesMut}; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; use crate::{ @@ -66,6 +69,7 @@ pub fn encode_raw_to_buf(frame: &RawFrame, buf: &mut Vec) -> ProtocolResult< }); } + buf.reserve(4 + frame_len as usize); buf.extend_from_slice(&frame_len.to_be_bytes()); buf.extend_from_slice(&frame.id.to_be_bytes()); buf.push(frame.flags); @@ -146,12 +150,13 @@ pub async fn read_raw_frame(reader: &mut R) -> ProtocolRes }); } - let mut payload = vec![0u8; frame_len]; - reader.read_exact(&mut payload).await?; - - let id = u32::from_be_bytes([payload[0], payload[1], payload[2], payload[3]]); - let flags = payload[4]; - let body = payload[FRAME_HEADER_SIZE..].to_vec(); + // Read the fixed header separately so the body is allocated exactly once. + let mut header = [0u8; FRAME_HEADER_SIZE]; + reader.read_exact(&mut header).await?; + let id = u32::from_be_bytes(header[..4].try_into().unwrap()); + let flags = header[4]; + let mut body = vec![0u8; frame_len - FRAME_HEADER_SIZE]; + reader.read_exact(&mut body).await?; Ok(RawFrame { id, flags, body }) } @@ -163,13 +168,60 @@ pub async fn write_raw_frame( writer: &mut W, frame: &RawFrame, ) -> ProtocolResult<()> { - let mut buf = Vec::new(); - encode_raw_to_buf(frame, &mut buf)?; - writer.write_all(&buf).await?; + let frame_len = u32::try_from(FRAME_HEADER_SIZE + frame.body.len()).map_err(|_| { + ProtocolError::FrameTooLarge { + size: u32::MAX, + max: MAX_FRAME_SIZE, + } + })?; + if frame_len > MAX_FRAME_SIZE { + return Err(ProtocolError::FrameTooLarge { + size: frame_len, + max: MAX_FRAME_SIZE, + }); + } + + let mut header = [0u8; 4 + FRAME_HEADER_SIZE]; + header[..4].copy_from_slice(&frame_len.to_be_bytes()); + header[4..8].copy_from_slice(&frame.id.to_be_bytes()); + header[8] = frame.flags; + write_vectored_all(writer, &header, &frame.body).await?; writer.flush().await?; Ok(()) } +/// Decode one complete typed frame from a cursor-based buffer without front-draining it. +pub fn try_decode_from_bytes(buf: &mut BytesMut) -> ProtocolResult> { + if buf.len() < 4 { + return Ok(None); + } + + let frame_len = u32::from_be_bytes(buf[..4].try_into().unwrap()); + if frame_len > MAX_FRAME_SIZE { + return Err(ProtocolError::FrameTooLarge { + size: frame_len, + max: MAX_FRAME_SIZE, + }); + } + + let frame_len = frame_len as usize; + if frame_len < FRAME_HEADER_SIZE { + return Err(ProtocolError::FrameTooShort { + size: frame_len as u32, + min: FRAME_HEADER_SIZE as u32, + }); + } + + let total = 4 + frame_len; + if buf.len() < total { + return Ok(None); + } + + let message = decode_message_frame(&buf[..total])?; + buf.advance(total); + Ok(Some(message)) +} + //-------------------------------------------------------------------------------------------------- // Functions: Typed message codec (CBOR-aware) //-------------------------------------------------------------------------------------------------- @@ -237,11 +289,17 @@ pub async fn write_message( writer: &mut W, message: &Message, ) -> ProtocolResult<()> { - let mut buf = Vec::new(); - encode_to_buf(message, &mut buf)?; - writer.write_all(&buf).await?; - writer.flush().await?; - Ok(()) + let mut body = Vec::new(); + ciborium::into_writer(message, &mut body)?; + write_raw_frame( + writer, + &RawFrame { + id: message.id, + flags: message.flags, + body, + }, + ) + .await } /// Decodes a [`RawFrame`] into a typed [`Message`] by CBOR-deserializing the body. @@ -288,12 +346,52 @@ pub fn decode_message_frame(frame: &[u8]) -> ProtocolResult { Ok(msg) } +async fn write_vectored_all( + writer: &mut W, + header: &[u8], + body: &[u8], +) -> std::io::Result<()> { + let mut header_offset = 0; + let mut body_offset = 0; + + while header_offset < header.len() || body_offset < body.len() { + let written = if header_offset < header.len() { + let slices = [ + IoSlice::new(&header[header_offset..]), + IoSlice::new(&body[body_offset..]), + ]; + let slice_count = if body_offset < body.len() { 2 } else { 1 }; + writer.write_vectored(&slices[..slice_count]).await? + } else { + // Some AsyncWrite implementations stop at an empty first IoSlice, so never leave the + // exhausted header in front of a non-empty body. + writer.write(&body[body_offset..]).await? + }; + if written == 0 { + return Err(std::io::ErrorKind::WriteZero.into()); + } + + let header_remaining = header.len() - header_offset; + if written < header_remaining { + header_offset += written; + } else { + header_offset = header.len(); + body_offset += written - header_remaining; + } + } + + Ok(()) +} + //-------------------------------------------------------------------------------------------------- // Tests //-------------------------------------------------------------------------------------------------- #[cfg(test)] mod tests { + use std::pin::Pin; + use std::task::{Context, Poll}; + use super::*; use crate::message::{FLAG_SESSION_START, FLAG_TERMINAL, MessageType, PROTOCOL_VERSION}; @@ -503,6 +601,46 @@ mod tests { assert_eq!(decoded.body, frame.body); } + #[test] + fn cursor_decoder_accepts_every_frame_fragment_boundary() { + let msg = Message::new(MessageType::Ready, 77, vec![0xAB; 1024]); + let mut encoded = Vec::new(); + encode_to_buf(&msg, &mut encoded).unwrap(); + + for split in 0..=encoded.len() { + let mut buf = BytesMut::new(); + buf.extend_from_slice(&encoded[..split]); + let first = try_decode_from_bytes(&mut buf).unwrap(); + if split < encoded.len() { + assert!(first.is_none(), "decoded incomplete frame at split {split}"); + buf.extend_from_slice(&encoded[split..]); + } + + let decoded = first + .or_else(|| try_decode_from_bytes(&mut buf).unwrap()) + .unwrap(); + assert_eq!(decoded.id, msg.id); + assert_eq!(decoded.t, msg.t); + assert!(buf.is_empty()); + } + } + + #[tokio::test] + async fn vectored_writer_handles_one_byte_short_writes() { + let frame = RawFrame { + id: 91, + flags: FLAG_TERMINAL, + body: vec![0xCD; 257], + }; + let mut expected = Vec::new(); + encode_raw_to_buf(&frame, &mut expected).unwrap(); + + let mut writer = OneByteWriter::default(); + write_raw_frame(&mut writer, &frame).await.unwrap(); + + assert_eq!(writer.bytes, expected); + } + #[test] fn test_raw_frame_sync_roundtrip() { let frame = RawFrame { @@ -539,4 +677,49 @@ mod tests { let payload: ExecExited = decoded.payload().unwrap(); assert_eq!(payload.code, 7); } + + #[derive(Default)] + struct OneByteWriter { + bytes: Vec, + } + + impl AsyncWrite for OneByteWriter { + fn poll_write( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + if let Some(byte) = buf.first() { + self.bytes.push(*byte); + Poll::Ready(Ok(1)) + } else { + Poll::Ready(Ok(0)) + } + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll> { + if let Some(byte) = bufs.iter().find_map(|buf| buf.first()) { + self.bytes.push(*byte); + Poll::Ready(Ok(1)) + } else { + Poll::Ready(Ok(0)) + } + } + } } diff --git a/crates/runtime/lib/clock.rs b/crates/runtime/lib/clock.rs index 1d0323751..387cc46ce 100644 --- a/crates/runtime/lib/clock.rs +++ b/crates/runtime/lib/clock.rs @@ -2,6 +2,7 @@ use std::time::{Duration, SystemTime}; +use bytes::Bytes; use microsandbox_protocol::codec; use microsandbox_protocol::core::ClockSync; use microsandbox_protocol::message::{Message, MessageType}; @@ -28,11 +29,11 @@ const CLOCK_SYNC_WAKE_THRESHOLD: Duration = Duration::from_secs(6); //-------------------------------------------------------------------------------------------------- /// Spawns a background task that keeps the guest wall clock aligned with the host. -pub(crate) fn spawn_clock_sync_task(agent_tx: mpsc::Sender>) -> JoinHandle<()> { +pub(crate) fn spawn_clock_sync_task(agent_tx: mpsc::Sender) -> JoinHandle<()> { tokio::spawn(clock_sync_task(agent_tx)) } -async fn clock_sync_task(agent_tx: mpsc::Sender>) { +async fn clock_sync_task(agent_tx: mpsc::Sender) { let mut last_wall = SystemTime::now(); let mut last_sync = match send_clock_sync(&agent_tx).await { Ok(sent_at) => sent_at, @@ -68,7 +69,7 @@ async fn clock_sync_task(agent_tx: mpsc::Sender>) { } } -async fn send_clock_sync(agent_tx: &mpsc::Sender>) -> RuntimeResult { +async fn send_clock_sync(agent_tx: &mpsc::Sender) -> RuntimeResult { let now = SystemTime::now(); let elapsed = now .duration_since(SystemTime::UNIX_EPOCH) @@ -84,7 +85,7 @@ async fn send_clock_sync(agent_tx: &mpsc::Sender>) -> RuntimeResult, + queued_bytes: Arc, + high_water_bytes: AtomicUsize, + full_events: AtomicU64, + byte_capacity: usize, +} + +/// One queue fragment whose charge follows unread bytes after it is popped. +pub struct QueuedBytes { + bytes: Bytes, + charge: ByteCharge, +} + +/// Releases byte capacity incrementally as a consumer advances its fragment cursor. +struct ByteCharge { + queued_bytes: Arc, + remaining: usize, +} + +/// Point-in-time observability for one console byte queue. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ByteQueueSnapshot { + /// Payload bytes currently retained. + pub queued_bytes: usize, + + /// Largest observed retained-byte count. + pub high_water_bytes: usize, + + /// Push attempts rejected by byte or entry capacity. + pub full_events: u64, + + /// Configured aggregate payload limit. + pub capacity: usize, +} + /// Shared state between the console port backend (libkrun threads) and the /// agent relay (tokio background tasks). /// @@ -55,16 +99,25 @@ const NAMED_PIPE_BRIDGE_TX_POLL_INTERVAL: Duration = Duration::from_millis(1); /// agent". pub struct ConsoleSharedState { /// Guest → Host: console TX thread pushes byte chunks, relay pops them. - pub tx_ring: ArrayQueue>, + pub tx_ring: ByteQueue, /// Host → Guest: relay pushes byte chunks, console RX thread pops them. - pub rx_ring: ArrayQueue>, + pub rx_ring: ByteQueue, /// Wakes the relay: "tx_ring has data from the guest." pub tx_wake: WakePipe, /// Wakes the console RX thread: "rx_ring has data for the guest." pub rx_wake: WakePipe, + + /// Wakes a blocked guest→host producer after the relay frees `tx_ring` capacity. + pub tx_capacity_wake: WakePipe, + + /// Wakes a blocked host→guest producer after libkrun frees `rx_ring` capacity. + pub rx_capacity_wake: WakePipe, + + /// Stops blocked console producers during teardown. + closed: AtomicBool, } /// Console port backend backed by [`ConsoleSharedState`]. @@ -81,7 +134,11 @@ pub struct AgentConsoleBackend { /// buffer. Protected by a Mutex because `read(&self)` takes `&self`. /// Only the RX thread calls `read`, so contention is zero. #[cfg(unix)] - pending: Mutex>, + pending: Mutex>, + + /// Size of the guest descriptor that most recently found `tx_ring` full. + #[cfg(unix)] + blocked_write_len: AtomicUsize, } #[cfg(windows)] @@ -96,20 +153,187 @@ pub(crate) struct AgentConsolePipeBridge { impl ConsoleSharedState { /// Create shared state with the default queue capacity. pub fn new() -> Self { - Self::with_capacity(DEFAULT_QUEUE_CAPACITY) + Self::with_capacity(DEFAULT_QUEUE_BYTE_CAPACITY) } - /// Create shared state with a specific queue capacity. - pub fn with_capacity(capacity: usize) -> Self { + /// Create shared state with a specific byte capacity in each direction. + pub fn with_capacity(byte_capacity: usize) -> Self { Self { - tx_ring: ArrayQueue::new(capacity), - rx_ring: ArrayQueue::new(capacity), + tx_ring: ByteQueue::new(byte_capacity), + rx_ring: ByteQueue::new(byte_capacity), tx_wake: WakePipe::new(), rx_wake: WakePipe::new(), + tx_capacity_wake: WakePipe::new(), + rx_capacity_wake: WakePipe::new(), + closed: AtomicBool::new(false), + } + } + + /// Unblock console producers because the runtime is shutting down. + pub fn close(&self) { + self.closed.store(true, Ordering::Release); + self.tx_capacity_wake.wake(); + self.rx_capacity_wake.wake(); + self.tx_wake.wake(); + self.rx_wake.wake(); + } + + /// Return whether no more console data should be admitted. + pub fn is_closed(&self) -> bool { + self.closed.load(Ordering::Acquire) + } +} + +impl ByteQueue { + /// Create a queue with an exact aggregate byte limit. + pub fn new(byte_capacity: usize) -> Self { + let entry_capacity = byte_capacity + .div_ceil(QUEUE_ENTRY_GRANULE) + .clamp(1, MAX_QUEUE_ENTRIES); + Self { + entries: ArrayQueue::new(entry_capacity), + queued_bytes: Arc::new(AtomicUsize::new(0)), + high_water_bytes: AtomicUsize::new(0), + full_events: AtomicU64::new(0), + byte_capacity, + } + } + + /// Push one owned byte region if both byte and entry capacity permit it. + pub fn push(&self, bytes: impl Into) -> Result<(), Bytes> { + let bytes = bytes.into(); + let len = bytes.len(); + + let reserved_bytes = loop { + let queued = self.queued_bytes.load(Ordering::Acquire); + let Some(next) = queued.checked_add(len) else { + self.full_events.fetch_add(1, Ordering::Relaxed); + return Err(bytes); + }; + if next > self.byte_capacity { + self.full_events.fetch_add(1, Ordering::Relaxed); + return Err(bytes); + } + if self + .queued_bytes + .compare_exchange_weak(queued, next, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + break next; + } + }; + self.high_water_bytes + .fetch_max(reserved_bytes, Ordering::Relaxed); + + let queued = QueuedBytes { + bytes, + charge: ByteCharge { + queued_bytes: Arc::clone(&self.queued_bytes), + remaining: len, + }, + }; + if let Err(queued) = self.entries.push(queued) { + self.full_events.fetch_add(1, Ordering::Relaxed); + return Err(queued.into_unqueued()); + } + Ok(()) + } + + /// Pop one fragment. Its unread-byte charge follows the returned owner. + pub fn pop(&self) -> Option { + self.entries.pop() + } + + /// Return whether a fragment of `len` bytes can be admitted now. + pub fn can_fit(&self, len: usize) -> bool { + self.queued_bytes + .load(Ordering::Acquire) + .checked_add(len) + .is_some_and(|next| next <= self.byte_capacity) + && !self.entries.is_full() + } + + /// Current queued payload bytes. + pub fn queued_bytes(&self) -> usize { + self.queued_bytes.load(Ordering::Acquire) + } + + /// Configured aggregate byte capacity. + pub fn capacity(&self) -> usize { + self.byte_capacity + } + + /// Capture bounded queue occupancy and saturation counters. + pub fn snapshot(&self) -> ByteQueueSnapshot { + ByteQueueSnapshot { + queued_bytes: self.queued_bytes(), + high_water_bytes: self.high_water_bytes.load(Ordering::Acquire), + full_events: self.full_events.load(Ordering::Relaxed), + capacity: self.capacity(), } } } +impl QueuedBytes { + /// Remaining unread bytes. + pub fn len(&self) -> usize { + self.bytes.len() + } + + /// Return whether the fragment cursor reached its end. + pub fn is_empty(&self) -> bool { + self.bytes.is_empty() + } + + /// Copy and release up to `out.len()` bytes from the front of the fragment. + #[cfg(unix)] + fn copy_prefix_into(&mut self, out: &mut [u8]) -> usize { + let len = self.len().min(out.len()); + out[..len].copy_from_slice(&self.bytes[..len]); + self.bytes.advance(len); + self.charge.release(len); + len + } + + /// Remove a failed queue admission while returning the caller's original bytes. + fn into_unqueued(mut self) -> Bytes { + let remaining = self.charge.remaining; + self.charge.release(remaining); + std::mem::take(&mut self.bytes) + } +} + +impl AsRef<[u8]> for QueuedBytes { + fn as_ref(&self) -> &[u8] { + &self.bytes + } +} + +impl std::ops::Deref for QueuedBytes { + type Target = [u8]; + + fn deref(&self) -> &Self::Target { + self.as_ref() + } +} + +impl ByteCharge { + fn release(&mut self, bytes: usize) { + debug_assert!(bytes <= self.remaining); + if bytes == 0 { + return; + } + self.remaining -= bytes; + self.queued_bytes.fetch_sub(bytes, Ordering::AcqRel); + } +} + +impl Drop for ByteCharge { + fn drop(&mut self) { + self.release(self.remaining); + } +} + impl AgentConsoleBackend { /// Create a new backend from shared state. pub fn new(shared: Arc) -> Self { @@ -117,7 +341,8 @@ impl AgentConsoleBackend { { Self { shared, - pending: Mutex::new(VecDeque::new()), + pending: Mutex::new(None), + blocked_write_len: AtomicUsize::new(0), } } @@ -187,28 +412,23 @@ impl ConsolePortBackend for AgentConsoleBackend { let mut pending = self.pending.lock().unwrap(); // Serve from leftover bytes first (use memcpy via slices). - if !pending.is_empty() { - let n = pending.len().min(buf.len()); - let (head, tail) = pending.as_slices(); - let from_head = n.min(head.len()); - buf[..from_head].copy_from_slice(&head[..from_head]); - if from_head < n { - let from_tail = n - from_head; - buf[from_head..n].copy_from_slice(&tail[..from_tail]); + if let Some(chunk) = pending.as_mut() { + let n = chunk.copy_prefix_into(buf); + if chunk.is_empty() { + pending.take(); } - pending.drain(..n); + self.shared.rx_capacity_wake.wake(); return Ok(n); } // Pop a new chunk from the ring. match self.shared.rx_ring.pop() { - Some(chunk) => { - let n = chunk.len().min(buf.len()); - buf[..n].copy_from_slice(&chunk[..n]); - // Buffer any remainder for subsequent reads. - if chunk.len() > buf.len() { - pending.extend(&chunk[buf.len()..]); + Some(mut chunk) => { + let n = chunk.copy_prefix_into(buf); + if !chunk.is_empty() { + *pending = Some(chunk); } + self.shared.rx_capacity_wake.wake(); Ok(n) } None => Err(io::ErrorKind::WouldBlock.into()), @@ -220,10 +440,18 @@ impl ConsolePortBackend for AgentConsoleBackend { /// Pushes a byte chunk to `tx_ring` and wakes the relay. Returns /// `WouldBlock` if the ring is full. fn write(&self, buf: &[u8]) -> io::Result { + if self.shared.is_closed() { + return Err(io::ErrorKind::BrokenPipe.into()); + } + self.shared .tx_ring - .push(buf.to_vec()) - .map_err(|_| io::Error::from(io::ErrorKind::WouldBlock))?; + .push(Bytes::copy_from_slice(buf)) + .map_err(|_| { + self.blocked_write_len.store(buf.len(), Ordering::Release); + io::Error::from(io::ErrorKind::WouldBlock) + })?; + self.blocked_write_len.store(0, Ordering::Release); self.shared.tx_wake.wake(); Ok(buf.len()) } @@ -233,6 +461,25 @@ impl ConsolePortBackend for AgentConsoleBackend { fn read_wake_fd(&self) -> RawFd { self.shared.rx_wake.as_raw_fd() } + + fn wait_until_writable(&self) { + loop { + let blocked_len = self.blocked_write_len.load(Ordering::Acquire).max(1); + if self.shared.is_closed() || self.shared.tx_ring.can_fit(blocked_len) { + return; + } + + // Drain then re-check before sleeping so a pop racing this transition cannot be lost. + self.shared.tx_capacity_wake.drain(); + if self.shared.is_closed() || self.shared.tx_ring.can_fit(blocked_len) { + return; + } + let _ = self + .shared + .tx_capacity_wake + .wait_timeout(Duration::from_secs(60)); + } + } } #[cfg(windows)] @@ -269,12 +516,15 @@ async fn bridge_guest_to_host( let mut buf = vec![0u8; NAMED_PIPE_BRIDGE_BUFFER_SIZE]; loop { + if shared.is_closed() { + return Ok(()); + } let n = reader.read(&mut buf).await?; if n == 0 { return Ok(()); } - push_queue_lossless(&shared.tx_ring, buf[..n].to_vec()).await; + push_queue_lossless(Arc::clone(&shared), Bytes::copy_from_slice(&buf[..n])).await; shared.tx_wake.wake(); } } @@ -285,9 +535,14 @@ async fn bridge_host_to_guest( shared: Arc, ) -> std::io::Result<()> { loop { + if shared.is_closed() { + return Ok(()); + } let mut wrote = false; while let Some(chunk) = shared.rx_ring.pop() { writer.write_all(&chunk).await?; + drop(chunk); + shared.rx_capacity_wake.wake(); wrote = true; } @@ -297,18 +552,42 @@ async fn bridge_host_to_guest( } shared.rx_wake.drain(); - tokio::time::sleep(NAMED_PIPE_BRIDGE_TX_POLL_INTERVAL).await; + if shared.rx_ring.queued_bytes() != 0 { + continue; + } + let shared_for_wait = Arc::clone(&shared); + let _ = tokio::task::spawn_blocking(move || { + shared_for_wait + .rx_wake + .wait_timeout(NAMED_PIPE_BRIDGE_WAIT_TIMEOUT) + }) + .await; } } #[cfg(windows)] -async fn push_queue_lossless(queue: &ArrayQueue>, mut chunk: Vec) { +async fn push_queue_lossless(shared: Arc, mut chunk: Bytes) { loop { - match queue.push(chunk) { + match shared.tx_ring.push(chunk) { Ok(()) => return, Err(returned) => { chunk = returned; - tokio::time::sleep(NAMED_PIPE_BRIDGE_TX_POLL_INTERVAL).await; + if shared.is_closed() { + return; + } + + // Drain then re-check before sleeping so a pop racing this transition cannot be lost. + shared.tx_capacity_wake.drain(); + if shared.tx_ring.can_fit(chunk.len()) { + continue; + } + let shared_for_wait = Arc::clone(&shared); + let _ = tokio::task::spawn_blocking(move || { + shared_for_wait + .tx_capacity_wake + .wait_timeout(NAMED_PIPE_BRIDGE_WAIT_TIMEOUT) + }) + .await; } } } @@ -333,7 +612,7 @@ mod tests { // Relay pops from tx_ring. let chunk = shared.tx_ring.pop().unwrap(); - assert_eq!(chunk, b"hello"); + assert_eq!(chunk.as_ref(), b"hello"); // Relay pushes response to rx_ring. shared.rx_ring.push(b"world".to_vec()).unwrap(); @@ -369,6 +648,54 @@ mod tests { assert_eq!(err.kind(), io::ErrorKind::WouldBlock); } + #[test] + fn byte_queue_releases_exact_capacity_on_pop() { + let queue = ByteQueue::new(8); + queue.push(Bytes::from_static(b"12345678")).unwrap(); + assert_eq!(queue.queued_bytes(), 8); + assert!(queue.push(Bytes::from_static(b"x")).is_err()); + + assert_eq!(queue.pop().unwrap().as_ref(), b"12345678"); + assert_eq!(queue.queued_bytes(), 0); + queue.push(Bytes::from_static(b"x")).unwrap(); + assert_eq!( + queue.snapshot(), + ByteQueueSnapshot { + queued_bytes: 1, + high_water_bytes: 8, + full_events: 1, + capacity: 8, + } + ); + } + + #[cfg(unix)] + #[test] + fn backend_capacity_wait_sleeps_until_consumer_pops() { + let shared = Arc::new(ConsoleSharedState::with_capacity(1)); + let backend = AgentConsoleBackend::new(Arc::clone(&shared)); + backend.write(b"a").unwrap(); + assert_eq!( + backend.write(b"b").unwrap_err().kind(), + io::ErrorKind::WouldBlock + ); + + let (done_tx, done_rx) = std::sync::mpsc::channel(); + let waiter = std::thread::spawn(move || { + backend.wait_until_writable(); + done_tx.send(()).unwrap(); + }); + assert!( + done_rx.recv_timeout(Duration::from_millis(25)).is_err(), + "waiter returned while the byte queue was still full" + ); + + shared.tx_ring.pop().unwrap(); + shared.tx_capacity_wake.wake(); + done_rx.recv_timeout(Duration::from_secs(1)).unwrap(); + waiter.join().unwrap(); + } + #[cfg(unix)] #[test] fn backend_read_drains_rx_wake_pipe() { @@ -416,7 +743,7 @@ mod tests { tokio::time::timeout(Duration::from_secs(1), async { loop { if let Some(bytes) = shared.tx_ring.pop() { - assert_eq!(bytes, b"guest-ready"); + assert_eq!(bytes.as_ref(), b"guest-ready"); return; } tokio::time::sleep(Duration::from_millis(1)).await; diff --git a/crates/runtime/lib/relay.rs b/crates/runtime/lib/relay.rs index f5f12dcf0..6cee3b784 100644 --- a/crates/runtime/lib/relay.rs +++ b/crates/runtime/lib/relay.rs @@ -9,12 +9,13 @@ //! handshake so that the relay can route agent responses back to the correct //! client without rewriting frame headers. -use std::collections::{HashMap, HashSet}; +use std::collections::{HashMap, HashSet, VecDeque}; +use std::io::IoSlice; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; -use bytes::{Bytes, BytesMut}; +use bytes::{Buf, Bytes, BytesMut}; #[cfg(unix)] use microsandbox_filesystem::{BindIdentityMap, BindIdentityMapHandle}; use microsandbox_protocol::codec::{self, MAX_FRAME_SIZE}; @@ -26,12 +27,12 @@ use microsandbox_protocol::message::{ use microsandbox_protocol::{AGENT_RELAY_ID_RANGE_STEP, AGENT_RELAY_MAX_CLIENTS}; #[cfg(unix)] use tokio::io::unix::AsyncFd; -use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt}; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; #[cfg(unix)] use tokio::net::UnixListener; #[cfg(windows)] use tokio::net::windows::named_pipe::{NamedPipeServer, PipeMode, ServerOptions}; -use tokio::sync::{Mutex, mpsc, watch}; +use tokio::sync::{Mutex, Semaphore, mpsc, watch}; use crate::clock::spawn_clock_sync_task; use crate::console::ConsoleSharedState; @@ -73,7 +74,23 @@ type SessionRegistry = std::sync::Mutex>; const LEN_PREFIX_SIZE: usize = 4; /// Capacity of the per-client write channel. -const CLIENT_WRITE_CHANNEL_CAPACITY: usize = 64; +const CLIENT_WRITE_CHANNEL_CAPACITY: usize = 2; + +/// Aggregate guest-to-client bytes retained by the relay. +const CLIENT_OUTPUT_BYTE_CAPACITY: usize = 32 * 1024 * 1024; + +/// Allocation granularity used for relay output admission. +const OUTPUT_BUDGET_GRANULE: usize = 4096; + +/// Maximum bytes opportunistically coalesced in one client socket batch. +const CLIENT_WRITE_BATCH_BYTES: usize = 256 * 1024; + +/// Maximum frame slices opportunistically coalesced in one client socket batch. +const CLIENT_WRITE_BATCH_FRAMES: usize = 64; + +/// At most eight generation-6 frames may wait between clients and the console. +/// Since a frame is capped at 4 MiB, this bounds the channel at 32 MiB. +const AGENT_WRITE_CHANNEL_CAPACITY: usize = 8; //-------------------------------------------------------------------------------------------------- // Types @@ -86,7 +103,13 @@ struct ClientState { /// Channel for sending frames to this client's writer task. /// Using a channel avoids holding the client mutex across async writes. /// Uses `Bytes` for zero-copy frame forwarding from the ring buffer. - write_tx: mpsc::Sender, + write_tx: mpsc::Sender, +} + +/// A client-bound frame whose aggregate capacity lives until the socket accepts it. +struct ClientWrite { + data: Bytes, + _permit: tokio::sync::OwnedSemaphorePermit, } /// The agent relay running in the sandbox process. @@ -322,6 +345,8 @@ impl AgentRelay { self.shared.tx_wake.drain(); while let Some(chunk) = self.shared.tx_ring.pop() { buf.extend_from_slice(&chunk); + drop(chunk); + self.shared.tx_capacity_wake.wake(); } // Try to extract complete frames. @@ -402,7 +427,7 @@ impl AgentRelay { mut shutdown: watch::Receiver, drain_tx: mpsc::Sender<()>, ) -> RuntimeResult<()> { - let ready_frame = self.ready_frame.ok_or_else(|| { + let ready_frame = self.ready_frame.take().ok_or_else(|| { RuntimeError::Custom("agent relay: run() called before wait_ready()".into()) })?; @@ -411,7 +436,7 @@ impl AgentRelay { // Bounded channel for client reader tasks to send frames to the ring writer. // Backpressure prevents unbounded memory growth from client floods. - let (agent_tx, agent_rx) = mpsc::channel::>(256); + let (agent_tx, agent_rx) = mpsc::channel::(AGENT_WRITE_CHANNEL_CAPACITY); // Track which client slots are in use. let used_slots: Arc>> = Arc::new(Mutex::new(HashSet::new())); @@ -440,11 +465,13 @@ impl AgentRelay { let next_session_id: Arc = Arc::new(AtomicU64::new(1)); let clients_for_reader = Arc::clone(&clients); let shared_for_reader = Arc::clone(&self.shared); + let client_output_budget = Arc::new(Semaphore::new(CLIENT_OUTPUT_BYTE_CAPACITY)); let log_writer_for_reader = self.log_writer.clone(); let registry_for_reader = Arc::clone(&session_registry); let ring_reader_handle = tokio::spawn(ring_reader_task( shared_for_reader, clients_for_reader, + client_output_budget, log_writer_for_reader, registry_for_reader, )); @@ -505,10 +532,37 @@ impl AgentRelay { // Spawn a per-client writer task so the ring reader // never holds the mutex across async writes. let (write_tx, mut write_rx) = - mpsc::channel::(CLIENT_WRITE_CHANNEL_CAPACITY); + mpsc::channel::(CLIENT_WRITE_CHANNEL_CAPACITY); tokio::spawn(async move { - while let Some(data) = write_rx.recv().await { - if let Err(e) = writer_half.write_all(&data).await { + let mut batch = VecDeque::new(); + let mut deferred = None; + loop { + let write = match deferred.take() { + Some(write) => write, + None => match write_rx.recv().await { + Some(write) => write, + None => break, + }, + }; + let mut batch_bytes = write.data.len(); + batch.push_back(write); + while batch.len() < CLIENT_WRITE_BATCH_FRAMES + && batch_bytes < CLIENT_WRITE_BATCH_BYTES + { + let Ok(write) = write_rx.try_recv() else { + break; + }; + if batch_bytes.saturating_add(write.data.len()) + > CLIENT_WRITE_BATCH_BYTES + { + deferred = Some(write); + break; + } + batch_bytes = batch_bytes.saturating_add(write.data.len()); + batch.push_back(write); + } + + if let Err(e) = write_client_batch(&mut writer_half, &mut batch).await { tracing::error!( "agent relay: client writer slot={slot} failed: {e}" ); @@ -568,6 +622,9 @@ impl AgentRelay { // Clean up the local IPC endpoint. self.listener.cleanup(&self.endpoint); + // Wake any libkrun or relay producer blocked on console capacity. + self.shared.close(); + // Abort background tasks. clock_sync_handle.abort(); ring_writer_handle.abort(); @@ -577,6 +634,28 @@ impl AgentRelay { } } +//-------------------------------------------------------------------------------------------------- +// Trait Implementations +//-------------------------------------------------------------------------------------------------- + +impl Drop for AgentRelay { + fn drop(&mut self) { + // The console write-capacity hook may be sleeping on a libkrun thread. Every relay exit, + // including readiness failure or task cancellation, must wake it before VM teardown. + self.shared.close(); + self.listener.cleanup(&self.endpoint); + let guest_to_host = self.shared.tx_ring.snapshot(); + let host_to_guest = self.shared.rx_ring.snapshot(); + tracing::debug!( + guest_to_host_high_water = guest_to_host.high_water_bytes, + guest_to_host_full_events = guest_to_host.full_events, + host_to_guest_high_water = host_to_guest.high_water_bytes, + host_to_guest_full_events = host_to_guest.full_events, + "agent relay console queue summary" + ); + } +} + //-------------------------------------------------------------------------------------------------- // Functions //-------------------------------------------------------------------------------------------------- @@ -590,10 +669,11 @@ pub(crate) fn push_guest_frame_blocking( pub(crate) fn push_guest_frame_until( shared: &ConsoleSharedState, - mut frame: Vec, + frame: Vec, timeout: std::time::Duration, ) -> RuntimeResult<()> { let deadline = std::time::Instant::now() + timeout; + let mut frame = Bytes::from(frame); loop { match shared.rx_ring.push(frame) { @@ -603,12 +683,19 @@ pub(crate) fn push_guest_frame_until( } Err(returned) => { frame = returned; - if std::time::Instant::now() >= deadline { + let remaining = deadline.saturating_duration_since(std::time::Instant::now()); + if remaining.is_zero() { return Err(RuntimeError::Custom( "timed out sending frame to agentd".into(), )); } - std::thread::sleep(std::time::Duration::from_millis(1)); + + // Drain then re-check to avoid losing a capacity transition racing the wait. + shared.rx_capacity_wake.drain(); + if shared.rx_ring.can_fit(frame.len()) { + continue; + } + let _ = shared.rx_capacity_wake.wait_timeout(remaining); } } } @@ -662,6 +749,37 @@ fn decode_frame(buf: &[u8]) -> RuntimeResult { codec::decode_message_frame(buf).map_err(|e| RuntimeError::Custom(format!("decode frame: {e}"))) } +/// Write a client batch with cursor advancement so short writes never compact frame tails. +async fn write_client_batch( + writer: &mut W, + batch: &mut VecDeque, +) -> std::io::Result<()> { + while !batch.is_empty() { + let slices: Vec> = batch + .iter() + .take(CLIENT_WRITE_BATCH_FRAMES) + .map(|write| IoSlice::new(&write.data)) + .collect(); + let written = writer.write_vectored(&slices).await?; + if written == 0 { + return Err(std::io::ErrorKind::WriteZero.into()); + } + + let mut remaining = written; + while remaining != 0 { + let front = batch.front_mut().expect("non-empty batch after write"); + if remaining < front.data.len() { + front.data.advance(remaining); + remaining = 0; + } else { + remaining -= front.data.len(); + batch.pop_front(); + } + } + } + writer.flush().await +} + /// Tap a guest-originated frame into `exec.log` if it belongs to the /// primary session. Best-effort: any decode error is logged and /// dropped — capture failures must never disrupt the routing path. @@ -727,7 +845,16 @@ fn tap_frame_into_log(frame: &RawFrame, writer: &LogWriter, session_registry: &S /// Background task that pushes client frames into the rx_ring for the guest. /// Retries on full ring with backoff to avoid dropping frames. -async fn ring_writer_task(shared: Arc, mut rx: mpsc::Receiver>) { +async fn ring_writer_task(shared: Arc, mut rx: mpsc::Receiver) { + #[cfg(unix)] + let capacity_fd = match AsyncFd::new(shared.rx_capacity_wake.as_raw_fd()) { + Ok(fd) => fd, + Err(error) => { + tracing::error!(%error, "agent relay: failed to watch console capacity"); + return; + } + }; + while let Some(frame_bytes) = rx.recv().await { let mut data = frame_bytes; let mut attempts = 0u64; @@ -746,7 +873,37 @@ async fn ring_writer_task(shared: Arc, mut rx: mpsc::Receive ); } data = returned; - tokio::time::sleep(std::time::Duration::from_millis(1)).await; + if shared.is_closed() { + return; + } + + shared.rx_capacity_wake.drain(); + if shared.rx_ring.can_fit(data.len()) { + continue; + } + + #[cfg(unix)] + { + let mut guard = match capacity_fd.readable().await { + Ok(guard) => guard, + Err(error) => { + tracing::error!(%error, "agent relay: console capacity wait failed"); + return; + } + }; + guard.clear_ready(); + } + + #[cfg(windows)] + { + let shared_for_wait = Arc::clone(&shared); + let _ = tokio::task::spawn_blocking(move || { + shared_for_wait + .rx_capacity_wake + .wait_timeout(std::time::Duration::from_secs(60)) + }) + .await; + } } } } @@ -765,6 +922,7 @@ async fn ring_writer_task(shared: Arc, mut rx: mpsc::Receive async fn ring_reader_task( shared: Arc, clients: Arc>>, + output_budget: Arc, log_writer: Option>, session_registry: Arc, ) { @@ -816,6 +974,8 @@ async fn ring_reader_task( shared.tx_wake.drain(); while let Some(chunk) = shared.tx_ring.pop() { buf.extend_from_slice(&chunk); + drop(chunk); + shared.tx_capacity_wake.wake(); } // Extract all complete frames first, then route them. @@ -854,7 +1014,26 @@ async fn ring_reader_task( match writer_result { Ok(write_tx) => { - if write_tx.send(frame.data).await.is_err() { + let charged = frame.data.len().saturating_add(OUTPUT_BUDGET_GRANULE - 1) + / OUTPUT_BUDGET_GRANULE + * OUTPUT_BUDGET_GRANULE; + let Ok(charged) = u32::try_from(charged) else { + tracing::error!("agent relay: client frame budget overflow"); + continue; + }; + let permit = match Arc::clone(&output_budget).acquire_many_owned(charged).await + { + Ok(permit) => permit, + Err(_) => return, + }; + if write_tx + .send(ClientWrite { + data: frame.data, + _permit: permit, + }) + .await + .is_err() + { tracing::error!("agent relay: write channel closed for slot={client_slot}"); } } @@ -930,7 +1109,7 @@ async fn read_raw_frame(reader: &mut R) -> RuntimeResul async fn client_reader_task( slot: u32, mut reader: impl AsyncRead + Unpin + Send + 'static, - agent_tx: mpsc::Sender>, + agent_tx: mpsc::Sender, clients: Arc>>, used_slots: Arc>>, drain_tx: mpsc::Sender<()>, @@ -1016,7 +1195,7 @@ async fn client_reader_task( } // Forward frame to ring writer (bounded — applies backpressure). - if agent_tx.send(frame.data.to_vec()).await.is_err() { + if agent_tx.send(frame.data).await.is_err() { tracing::error!("agent relay: ring writer channel closed"); break; } @@ -1061,7 +1240,7 @@ async fn client_reader_task( continue; } - if agent_tx.send(buf).await.is_err() { + if agent_tx.send(Bytes::from(buf)).await.is_err() { tracing::error!("agent relay: ring writer channel closed during cleanup"); break; } @@ -1085,7 +1264,7 @@ async fn client_reader_task( let mut buf = Vec::new(); match codec::encode_to_buf(&disconnect_msg, &mut buf) { Ok(()) => { - if agent_tx.send(buf).await.is_err() { + if agent_tx.send(Bytes::from(buf)).await.is_err() { tracing::error!("agent relay: ring writer channel closed during fs cleanup"); } } @@ -1116,6 +1295,9 @@ fn is_client_frame_allowed(id: u32, flags: u8, id_start: u32, id_end_exclusive: #[cfg(test)] mod tests { + use std::pin::Pin; + use std::task::{Context, Poll}; + use super::*; use microsandbox_protocol::core::Ready; @@ -1173,10 +1355,37 @@ mod tests { assert!(is_client_frame_allowed(0, FLAG_SHUTDOWN, 10, 20)); } + #[tokio::test] + async fn client_batch_handles_short_vectored_writes_and_releases_budget() { + let budget = Arc::new(Semaphore::new(8)); + let first = Arc::clone(&budget).acquire_many_owned(3).await.unwrap(); + let second = Arc::clone(&budget).acquire_many_owned(5).await.unwrap(); + let mut batch = VecDeque::from([ + ClientWrite { + data: Bytes::from_static(b"abc"), + _permit: first, + }, + ClientWrite { + data: Bytes::from_static(b"defgh"), + _permit: second, + }, + ]); + let mut writer = ShortVectoredWriter { + max_write: 2, + ..Default::default() + }; + + write_client_batch(&mut writer, &mut batch).await.unwrap(); + + assert_eq!(writer.bytes, b"abcdefgh"); + assert!(batch.is_empty()); + assert_eq!(budget.available_permits(), 8); + } + #[tokio::test] #[cfg(unix)] async fn wait_ready_rejects_ready_before_init_when_maps_are_pending() { - let shared = Arc::new(ConsoleSharedState::with_capacity(8)); + let shared = Arc::new(ConsoleSharedState::with_capacity(64 * 1024)); let handle = Arc::new(std::sync::OnceLock::new()); let sock_path = test_agent_endpoint("ready-before-init"); let mut relay = AgentRelay::new(&sock_path, Arc::clone(&shared)) @@ -1209,7 +1418,7 @@ mod tests { #[tokio::test] #[cfg(unix)] async fn wait_ready_installs_init_map_before_ready() { - let shared = Arc::new(ConsoleSharedState::with_capacity(8)); + let shared = Arc::new(ConsoleSharedState::with_capacity(64 * 1024)); let handle = Arc::new(std::sync::OnceLock::new()); let sock_path = test_agent_endpoint("init-map"); let mut relay = AgentRelay::new(&sock_path, Arc::clone(&shared)) @@ -1258,7 +1467,7 @@ mod tests { #[tokio::test] async fn wait_ready_skips_init_requirement_when_no_bind_map_pending() { - let shared = Arc::new(ConsoleSharedState::with_capacity(8)); + let shared = Arc::new(ConsoleSharedState::with_capacity(64 * 1024)); let sock_path = test_agent_endpoint("no-bind-map"); let mut relay = AgentRelay::new(&sock_path, Arc::clone(&shared)) .await @@ -1288,4 +1497,53 @@ mod tests { "no init context means no ack should be sent" ); } + + #[derive(Default)] + struct ShortVectoredWriter { + bytes: Vec, + max_write: usize, + } + + impl AsyncWrite for ShortVectoredWriter { + fn poll_write( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let len = buf.len().min(self.max_write); + self.bytes.extend_from_slice(&buf[..len]); + Poll::Ready(Ok(len)) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + bufs: &[IoSlice<'_>], + ) -> Poll> { + let mut remaining = self.max_write; + let mut written = 0; + for buf in bufs { + let len = buf.len().min(remaining); + self.bytes.extend_from_slice(&buf[..len]); + written += len; + remaining -= len; + if remaining == 0 { + break; + } + } + Poll::Ready(Ok(written)) + } + } } diff --git a/crates/runtime/lib/vm.rs b/crates/runtime/lib/vm.rs index f9e1a08b0..80c336e71 100644 --- a/crates/runtime/lib/vm.rs +++ b/crates/runtime/lib/vm.rs @@ -2931,7 +2931,7 @@ mod tests { request_guest_shutdown(&shared).unwrap(); - let mut frame = shared.rx_ring.pop().unwrap(); + let mut frame = shared.rx_ring.pop().unwrap().to_vec(); let msg = codec::try_decode_from_buf(&mut frame).unwrap().unwrap(); assert_eq!(msg.t, MessageType::Shutdown); assert_eq!(msg.id, 0); @@ -2975,7 +2975,7 @@ mod tests { #[test] fn test_request_guest_shutdown_with_timeout_fails_when_ring_full() { - let shared = ConsoleSharedState::with_capacity(1); + let shared = ConsoleSharedState::with_capacity(8); shared.rx_ring.push(b"occupied".to_vec()).unwrap(); let err = request_guest_shutdown_with_timeout(&shared, Duration::ZERO).unwrap_err(); diff --git a/packages/agent-client/rust/lib/client.rs b/packages/agent-client/rust/lib/client.rs index 09f4f4afb..6e288c5e2 100644 --- a/packages/agent-client/rust/lib/client.rs +++ b/packages/agent-client/rust/lib/client.rs @@ -66,9 +66,11 @@ const DEFAULT_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10); const WINDOWS_PIPE_CONNECT_RETRY: Duration = Duration::from_millis(10); #[cfg(feature = "stream")] -const WRITER_QUEUE_CAPACITY: usize = 1024; +/// Eight maximum-sized generation-6 frames bound queued writes at 32 MiB. +const WRITER_QUEUE_CAPACITY: usize = 8; const REQUEST_QUEUE_CAPACITY: usize = 1; -const STREAM_QUEUE_CAPACITY: usize = 1024; +/// Two maximum-sized frames keep each correlation stream at or below 8 MiB. +const STREAM_QUEUE_CAPACITY: usize = 2; const LEGACY_PROTOCOL_VERSION: u8 = 1; // TODO(upgrade-0.6): Remove in 0.6.x or later once live-sandbox diff --git a/sdk/rust/lib/sandbox/fs.rs b/sdk/rust/lib/sandbox/fs.rs index 2170f57f6..5357694c0 100644 --- a/sdk/rust/lib/sandbox/fs.rs +++ b/sdk/rust/lib/sandbox/fs.rs @@ -126,6 +126,8 @@ pub struct FsReadStream { // `fs_read_stream` returns and `rx` would receive nothing. client: Option>, close_handle: Option, + /// Set only after a terminal `FsResponse`; channel closure alone is an error. + finished: bool, } /// A streaming writer for file data to the sandbox. @@ -535,6 +537,7 @@ impl FsReadStream { rx, client: Some(client), close_handle, + finished: false, } } @@ -543,6 +546,10 @@ impl FsReadStream { /// Returns `None` when the stream is complete (after `FsResponse`). /// Returns an error if the guest reported a failure. pub async fn recv(&mut self) -> MicrosandboxResult> { + if self.finished { + return Ok(None); + } + while let Some(msg) = self.rx.recv().await { match msg.t { MessageType::FsData => { @@ -554,6 +561,7 @@ impl FsReadStream { MessageType::FsResponse => { let resp: FsResponse = msg.payload()?; let close_result = self.close_owned_handle().await; + self.finished = true; if !resp.ok { return Err(MicrosandboxError::SandboxFsOps( resp.error.unwrap_or_else(|| "unknown error".into()), @@ -566,7 +574,9 @@ impl FsReadStream { } } self.close_owned_handle().await?; - Ok(None) + Err(MicrosandboxError::SandboxFsOps( + "filesystem read stream closed before terminal response".into(), + )) } /// Collect all remaining data into bytes. @@ -608,9 +618,12 @@ impl FsWriteSink { /// Write a chunk of data. pub async fn write(&self, data: impl AsRef<[u8]>) -> MicrosandboxResult<()> { - let fs_data = FsData { - data: data.as_ref().to_vec(), - }; + self.write_owned(data.as_ref().to_vec()).await + } + + /// Write an already-owned chunk without cloning it at the SDK stream boundary. + async fn write_owned(&self, data: Vec) -> MicrosandboxResult<()> { + let fs_data = FsData { data }; self.client .send(self.id, MessageType::FsData, &fs_data) .await @@ -748,7 +761,7 @@ pub(crate) mod agent { }, message::MessageType, }; - use tokio::io::AsyncReadExt; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; use crate::{MicrosandboxError, MicrosandboxResult, agent::AgentClient, backend::Backend}; @@ -1347,13 +1360,14 @@ pub(crate) mod agent { ) -> MicrosandboxResult<()> { let mut file = tokio::fs::File::open(host_path).await?; let sink = write_stream(backend, name, guest_path).await?; - let mut buf = vec![0u8; FS_CHUNK_SIZE]; loop { + let mut buf = vec![0u8; FS_CHUNK_SIZE]; let n = file.read(&mut buf).await?; if n == 0 { break; } - sink.write(&buf[..n]).await?; + buf.truncate(n); + sink.write_owned(buf).await?; } sink.close().await } @@ -1364,10 +1378,303 @@ pub(crate) mod agent { guest_path: &str, host_path: &Path, ) -> MicrosandboxResult<()> { - let data = read(backend, name, guest_path).await?; - tokio::fs::write(host_path, &data).await?; + let (std_file, temp_path) = prepare_host_copy_target(host_path).await?; + let mut file = tokio::fs::File::from_std(std_file); + let mut stream = read_stream(backend, name, guest_path).await?; + let mut received = 0u64; + + while let Some(chunk) = stream.recv().await? { + file.write_all(&chunk).await?; + received = received.checked_add(chunk.len() as u64).ok_or_else(|| { + MicrosandboxError::SandboxFsOps( + "copied file size exceeds the supported u64 range".into(), + ) + })?; + } + + file.flush().await?; + let written = file.metadata().await?.len(); + if written != received { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("host copy byte-count mismatch: received {received}, wrote {written}"), + ) + .into()); + } + file.sync_all().await?; + drop(file); + + publish_host_copy_target(temp_path, host_path)?; + tracing::debug!(bytes = received, path = %host_path.display(), "copied guest file to host"); + Ok(()) + } + + async fn prepare_host_copy_target( + host_path: &Path, + ) -> MicrosandboxResult<(std::fs::File, tempfile::TempPath)> { + let host_path = host_path.to_path_buf(); + tokio::task::spawn_blocking(move || { + let existing_permissions = match std::fs::symlink_metadata(&host_path) { + Ok(metadata) if metadata.file_type().is_symlink() => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "refusing to replace symbolic-link destination {}", + host_path.display() + ), + )); + } + Ok(metadata) if metadata.is_dir() => { + return Err(std::io::Error::new( + std::io::ErrorKind::IsADirectory, + format!("copy destination is a directory: {}", host_path.display()), + )); + } + Ok(metadata) => Some(metadata.permissions()), + Err(error) if error.kind() == std::io::ErrorKind::NotFound => None, + Err(error) => return Err(error), + }; + + let parent = host_path + .parent() + .filter(|path| !path.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + let file_name = host_path.file_name().ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("copy destination has no file name: {}", host_path.display()), + ) + })?; + let prefix = format!(".{}.msb-copy-", file_name.to_string_lossy()); + let named = tempfile::Builder::new() + .prefix(&prefix) + .tempfile_in(parent)?; + if let Some(permissions) = existing_permissions { + named.as_file().set_permissions(permissions)?; + } else { + // NamedTempFile is owner-only by default. Set the mode explicitly so this + // security property does not depend on a future tempfile implementation. + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + + named + .as_file() + .set_permissions(std::fs::Permissions::from_mode(0o600))?; + } + } + let (file, path) = named.into_parts(); + Ok((file, path)) + }) + .await + .map_err(|error| MicrosandboxError::Custom(format!("host copy worker failed: {error}")))? + .map_err(Into::into) + } + + fn publish_host_copy_target( + temp_path: tempfile::TempPath, + host_path: &Path, + ) -> MicrosandboxResult<()> { + // Re-check immediately before rename. Atomic rename never follows a symlink, but + // rejecting it keeps the public behavior explicit even if the path changed mid-copy. + match std::fs::symlink_metadata(host_path) { + Ok(metadata) if metadata.file_type().is_symlink() => { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!( + "refusing to replace symbolic-link destination {}", + host_path.display() + ), + ) + .into()); + } + Ok(_) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(error.into()), + } + + // Keep the final lstat+rename in one non-awaiting region. Once publication starts, task + // cancellation cannot report failure while a detached blocking worker commits later. + temp_path.persist(host_path).map_err(|error| error.error)?; Ok(()) } + + #[cfg(test)] + mod tests { + #[cfg(unix)] + use std::os::unix::fs::{PermissionsExt, symlink}; + + use super::*; + + #[tokio::test] + async fn host_copy_target_atomically_replaces_existing_file() { + let dir = tempfile::tempdir().unwrap(); + let destination = dir.path().join("artifact.bin"); + std::fs::write(&destination, b"old").unwrap(); + + let (file, temp_path) = prepare_host_copy_target(&destination).await.unwrap(); + let mut file = tokio::fs::File::from_std(file); + file.write_all(b"complete replacement").await.unwrap(); + file.sync_all().await.unwrap(); + drop(file); + publish_host_copy_target(temp_path, &destination).unwrap(); + + assert_eq!( + std::fs::read(&destination).unwrap(), + b"complete replacement" + ); + } + + #[cfg(unix)] + #[tokio::test] + async fn host_copy_target_preserves_unix_mode() { + let dir = tempfile::tempdir().unwrap(); + let destination = dir.path().join("artifact.bin"); + std::fs::write(&destination, b"old").unwrap(); + std::fs::set_permissions(&destination, std::fs::Permissions::from_mode(0o640)).unwrap(); + + let (file, temp_path) = prepare_host_copy_target(&destination).await.unwrap(); + drop(file); + publish_host_copy_target(temp_path, &destination).unwrap(); + + assert_eq!( + std::fs::metadata(&destination) + .unwrap() + .permissions() + .mode() + & 0o777, + 0o640 + ); + } + + #[tokio::test] + async fn cancelled_host_copy_removes_temp_and_keeps_destination() { + let dir = tempfile::tempdir().unwrap(); + let destination = dir.path().join("artifact.bin"); + std::fs::write(&destination, b"original").unwrap(); + + let (file, temp_path) = prepare_host_copy_target(&destination).await.unwrap(); + let temp_name = temp_path.to_path_buf(); + drop(file); + drop(temp_path); + + assert!(!temp_name.exists()); + assert_eq!(std::fs::read(&destination).unwrap(), b"original"); + } + + #[cfg(unix)] + #[tokio::test] + async fn new_host_copy_target_is_owner_only() { + let dir = tempfile::tempdir().unwrap(); + let destination = dir.path().join("new.bin"); + let (file, temp_path) = prepare_host_copy_target(&destination).await.unwrap(); + drop(file); + publish_host_copy_target(temp_path, &destination).unwrap(); + + assert_eq!( + std::fs::metadata(&destination) + .unwrap() + .permissions() + .mode() + & 0o777, + 0o600 + ); + } + + #[cfg(unix)] + #[tokio::test] + async fn host_copy_rejects_symbolic_link_destination() { + let dir = tempfile::tempdir().unwrap(); + let target = dir.path().join("target.bin"); + let destination = dir.path().join("link.bin"); + std::fs::write(&target, b"target").unwrap(); + symlink(&target, &destination).unwrap(); + + let error = prepare_host_copy_target(&destination).await.unwrap_err(); + assert!(matches!( + error, + MicrosandboxError::Io(ref error) + if error.kind() == std::io::ErrorKind::InvalidInput + )); + assert_eq!(std::fs::read(&target).unwrap(), b"target"); + } + + #[cfg(unix)] + #[tokio::test] + async fn host_copy_rechecks_symbolic_link_before_publish() { + let dir = tempfile::tempdir().unwrap(); + let target = dir.path().join("target.bin"); + let destination = dir.path().join("link.bin"); + std::fs::write(&target, b"target").unwrap(); + let (file, temp_path) = prepare_host_copy_target(&destination).await.unwrap(); + drop(file); + symlink(&target, &destination).unwrap(); + + let error = publish_host_copy_target(temp_path, &destination).unwrap_err(); + assert!(matches!( + error, + MicrosandboxError::Io(ref error) + if error.kind() == std::io::ErrorKind::InvalidInput + )); + assert!( + std::fs::symlink_metadata(&destination) + .unwrap() + .file_type() + .is_symlink() + ); + } + } +} + +//-------------------------------------------------------------------------------------------------- +// Tests +//-------------------------------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn read_stream_rejects_channel_close_without_terminal_response() { + let (tx, rx) = mpsc::channel(1); + drop(tx); + let mut stream = FsReadStream { + rx, + client: None, + close_handle: None, + finished: false, + }; + + let error = stream.recv().await.unwrap_err(); + assert!( + error + .to_string() + .contains("closed before terminal response") + ); + } + + #[tokio::test] + async fn read_stream_finishes_only_after_success_response() { + let (tx, rx) = mpsc::channel(1); + let response = FsResponse { + ok: true, + error: None, + data: None, + }; + tx.send(Message::with_payload(MessageType::FsResponse, 1, &response).unwrap()) + .await + .unwrap(); + drop(tx); + let mut stream = FsReadStream { + rx, + client: None, + close_handle: None, + finished: false, + }; + + assert!(stream.recv().await.unwrap().is_none()); + assert!(stream.recv().await.unwrap().is_none()); + } } //--------------------------------------------------------------------------------------------------