1use std::collections::{HashMap, HashSet, VecDeque};
4use std::env;
5use std::fs::{File, OpenOptions};
6use std::os::fd::AsRawFd;
7use std::sync::Arc;
8use std::sync::atomic::{AtomicBool, Ordering};
9use std::task::{Context, Poll};
10use std::time::Instant;
11
12use bytes::BytesMut;
13use chrono::Utc;
14use tokio::io::unix::AsyncFd;
15use tokio::sync::{OwnedSemaphorePermit, Semaphore, watch};
16use tokio::time::{self, Duration};
17
18use microsandbox_protocol::AGENT_TRANSPORT_DUAL_PORT_CMDLINE;
19use microsandbox_protocol::bootstrap::GuestBootstrap;
20use microsandbox_protocol::bulk::{
21 BULK_HEADER_SIZE, BULK_PROTOCOL_VERSION, BulkCancel, BulkCancelReason, BulkCredit, BulkFinish,
22 BulkFlow, BulkKind, BulkRecord, DEFAULT_BULK_WINDOW, DEFAULT_FILESYSTEM_BULK_RECORD_PAYLOAD,
23 MAX_BULK_RECORD_PAYLOAD, MIN_BULK_RECORD_PAYLOAD,
24};
25use microsandbox_protocol::codec::{self, DecodedFrame, MAX_FRAME_SIZE};
26use microsandbox_protocol::core::{
27 ClockSync, CoreError, CoreErrorKind, InitAck, InitFailureReason, InitResolved, Ping, Pong,
28 Ready, RelayClientDisconnected, ResolvedUser, Touch, Touched,
29 WORKLOAD_TRANSPORT_BARRIER_VERSION, WORKLOAD_TRANSPORT_BULK_BYTES,
30 WORKLOAD_TRANSPORT_BULK_FRAMES, WORKLOAD_TRANSPORT_CONTROL_BYTES,
31 WORKLOAD_TRANSPORT_CONTROL_FRAMES, WorkloadFailure, WorkloadFailureDisposition, WorkloadFreeze,
32 WorkloadFrozen, WorkloadThaw, WorkloadThawed, WorkloadTransportCredit,
33 WorkloadTransportPosition,
34};
35use microsandbox_protocol::exec::{
36 ExecExited, ExecFailed, ExecFailureKind, ExecRequest, ExecResize, ExecSignal, ExecStarted,
37 ExecStderr, ExecStdin, ExecStdinError, ExecStdout,
38};
39use microsandbox_protocol::fs::{FS_CHUNK_SIZE, FsData, FsRequest, FsResponse};
40use microsandbox_protocol::heartbeat::{ActivityCounters, Heartbeat};
41use microsandbox_protocol::message::{FRAME_HEADER_SIZE, Message, MessageType};
42use microsandbox_protocol::tcp::{TcpClose, TcpConnect, TcpData, TcpEof, TcpFailed};
43use microsandbox_protocol::transport::{
44 BULK_BINDING_SIZE, BulkTransportReady, CLIENT_INCARNATION_SIZE, ClientIncarnation,
45 IncarnatedBulkFrame, RelayClientConnected, RelayClientDisconnectedAck, RelayLeaseReady,
46 decode_bulk_ack, encode_bulk_hello, encode_relay_client_disconnected_ack,
47 relay_client_id_range, relay_client_slot, try_decode_incarnated_bulk_from_bytes,
48 try_decode_relay_client_connected_from_bytes,
49 try_decode_relay_client_disconnected_ack_from_bytes,
50 validate_relay_client_range as canonical_relay_client_range,
51};
52
53use crate::config::{AgentdConfig, scripts_path};
54use crate::error::{AgentdError, AgentdResult};
55use crate::fs::{FsReadSession, FsState, FsStreamSession, FsWriteSession};
56use crate::process::ProcessManager;
57use crate::serial::{AGENT_BULK_PORT_NAME, AGENT_PORT_NAME, InputCharge, InputLane, InputWindow};
58use crate::session::{
59 BulkOutputCommand, ExecSession, RawActivity, RawSessionCompletion, RawSessionOutput,
60 SessionOutput, SessionOutputEnvelope, SessionOutputSender, resolve_default_user,
61};
62use crate::tcp::TcpSession;
63use crate::workload::{WorkloadLatch, WorkloadLatchError};
64use crate::{clock, fs, handoff, heartbeat, serial};
65
66const HEARTBEAT_INTERVAL_SECS: u64 = 1;
75
76const CONTROL_SERIAL_READ_BUF_SIZE: usize = 64 * 1024;
79
80const COMBINED_SERIAL_READ_BUF_SIZE: usize = DEFAULT_FILESYSTEM_BULK_RECORD_PAYLOAD as usize + 32;
82
83const BULK_SERIAL_READ_BUF_SIZE: usize =
85 DEFAULT_FILESYSTEM_BULK_RECORD_PAYLOAD as usize + CLIENT_INCARNATION_SIZE + 32;
86
87const MAX_INPUT_BUF_SIZE: usize = MAX_FRAME_SIZE as usize + 4;
89
90const MAX_BULK_INPUT_BUF_SIZE: usize = MAX_INPUT_BUF_SIZE + CLIENT_INCARNATION_SIZE;
92
93const AGENT_READ_QUANTUM_BYTES: usize = 256 * 1024;
96const AGENT_READ_QUANTUM_CALLS: usize = 64;
97
98const INIT_ACK_TIMEOUT_SECS: u64 = 60;
100
101const BULK_BINDING_TIMEOUT_SECS: u64 = 60;
103
104const BULK_SCHEDULER_QUANTUM: usize = 256 * 1024;
106
107const BULK_SCHEDULER_MAX_BURST: usize = MAX_BULK_RECORD_PAYLOAD as usize;
109
110const BULK_SCHEDULER_FLOW_CAPACITY: usize = 8 * 1024 * 1024;
112
113const BULK_SCHEDULER_MAX_FLOWS_PER_CLIENT: usize = 64;
115
116const BULK_READER_MAX_RECORDS_PER_TURN: usize = 64;
118
119const BULK_READER_MAX_BYTES_PER_TURN: usize = 1024 * 1024;
121
122const FS_BULK_INPUT_ITEM_CAPACITY: usize =
124 DEFAULT_BULK_WINDOW as usize / MIN_BULK_RECORD_PAYLOAD as usize;
125
126const BULK_INPUT_BYTE_CAPACITY: usize = 32 * 1024 * 1024;
128
129const FS_ACTIVITY_BATCH_BYTES: usize = 4 * 1024 * 1024;
131
132const FS_BULK_WRITE_COALESCE_BYTES: usize = FS_CHUNK_SIZE;
134
135const BULK_ACTIVITY_PUBLISH_BYTES: usize = 4 * 1024 * 1024;
137const BULK_ACTIVITY_PUBLISH_INTERVAL: Duration = Duration::from_millis(100);
138
139const BULK_FAILURE_FLUSH_TIMEOUT: Duration = Duration::from_secs(2);
141
142struct AgentState {
147 input_window: InputWindow,
148 pending_freeze: Option<Message>,
149 aborted_transport_attempt: Option<String>,
150 output_parked: bool,
151 resume_output_after_flush: bool,
152 frozen_host_input: Option<WorkloadTransportPosition>,
153 frozen_guest_bulk_bytes: u64,
154 detached_sessions: HashMap<(u64, u32), ExecSession>,
156 stdin_poll_offset: usize,
157 restored_attempt: Option<String>,
158 client_incarnations: HashMap<u32, ClientIncarnation>,
159 bulk_input_budget: Arc<Semaphore>,
160 sessions: HashMap<u32, ExecSession>,
161 write_sessions: HashMap<u32, FsWriteSession>,
162 bulk_write_workers: HashMap<u32, FsBulkWriteWorker>,
163 read_sessions: HashMap<u32, FsReadSession>,
164 tcp_sessions: HashMap<u32, TcpSession>,
165 bulk_received_offsets: HashMap<u32, u64>,
166 pending_bulk_finishes: HashMap<u32, BulkFinish>,
167 fs: FsState,
168}
169
170struct FsBulkWriteWorker {
172 records: tokio::sync::mpsc::Sender<AdmittedBulkRecord>,
173 finish: tokio::sync::mpsc::Sender<BulkFinish>,
174 task: tokio::task::JoinHandle<()>,
175}
176
177pub(crate) struct AdmittedBulkRecord {
180 record: BulkRecord,
181 _permit: OwnedSemaphorePermit,
182 _transport_charge: Option<InputCharge>,
183}
184
185pub(crate) struct BulkInputPermit {
187 _payload: OwnedSemaphorePermit,
188 _transport: Option<InputCharge>,
189}
190
191struct PendingBulkInput {
193 frame: IncarnatedBulkFrame,
194 budget: Arc<Semaphore>,
195 charge: InputCharge,
196}
197
198struct ActivityTracker {
199 activity_seq: u64,
200 counters: ActivityCounters,
201}
202
203#[derive(Default)]
208pub struct BootConsoleState {
209 input: Vec<u8>,
210 init_failure_ack: bool,
211}
212
213#[derive(Clone)]
214struct HeartbeatSnapshot {
215 activity_seq: u64,
216 active_exec_sessions: u32,
217 active_fs_streams: u32,
218 active_tcp_streams: u32,
219 counters: ActivityCounters,
220}
221
222pub struct BoundBulkPort {
224 file: File,
225 connection_id: [u8; 16],
226}
227
228struct BulkInputState {
233 port: AsyncFd<File>,
234 read_buf: Vec<u8>,
235 input: BytesMut,
236}
237
238#[derive(Default)]
240struct AgentReadBudget {
241 bytes: usize,
242 calls: usize,
243}
244
245struct BulkWriteFlow {
247 queue: VecDeque<SessionOutputEnvelope>,
248 queued_bytes: usize,
249 deficit: usize,
250}
251
252enum BulkOutputCleanup {
254 Park {
255 completion: tokio::sync::oneshot::Sender<u64>,
256 wire_bytes: u64,
257 },
258 Flow(tokio::sync::oneshot::Sender<()>),
259 Incarnation {
260 incarnation: ClientIncarnation,
261 completion: tokio::sync::oneshot::Sender<()>,
262 },
263}
264
265#[derive(Default)]
266struct BulkOutputPosition {
267 parked: bool,
268 wire_bytes: u64,
269}
270
271impl Default for AgentState {
276 fn default() -> Self {
277 Self {
278 input_window: InputWindow::new(initial_input_credit()),
279 pending_freeze: None,
280 aborted_transport_attempt: None,
281 output_parked: false,
282 resume_output_after_flush: false,
283 frozen_host_input: None,
284 frozen_guest_bulk_bytes: 0,
285 detached_sessions: HashMap::new(),
286 stdin_poll_offset: 0,
287 restored_attempt: None,
288 client_incarnations: HashMap::new(),
289 bulk_input_budget: Arc::new(Semaphore::new(BULK_INPUT_BYTE_CAPACITY)),
290 sessions: HashMap::new(),
291 write_sessions: HashMap::new(),
292 bulk_write_workers: HashMap::new(),
293 read_sessions: HashMap::new(),
294 tcp_sessions: HashMap::new(),
295 bulk_received_offsets: HashMap::new(),
296 pending_bulk_finishes: HashMap::new(),
297 fs: FsState::default(),
298 }
299 }
300}
301
302impl AdmittedBulkRecord {
303 pub(crate) fn record(&self) -> &BulkRecord {
304 &self.record
305 }
306
307 pub(crate) fn into_parts(self) -> (BulkRecord, BulkInputPermit) {
308 (
309 self.record,
310 BulkInputPermit {
311 _payload: self._permit,
312 _transport: self._transport_charge,
313 },
314 )
315 }
316
317 #[cfg(test)]
318 pub(crate) fn for_test(record: BulkRecord) -> Self {
319 let payload_len = record.payload.len();
320 let budget = Arc::new(Semaphore::new(payload_len));
321 let permit = budget
322 .try_acquire_many_owned(payload_len as u32)
323 .expect("test record fits its exact input budget");
324 Self {
325 record,
326 _permit: permit,
327 _transport_charge: None,
328 }
329 }
330}
331
332impl ActivityTracker {
333 fn new() -> Self {
334 Self {
335 activity_seq: 0,
336 counters: ActivityCounters::default(),
337 }
338 }
339
340 fn record_host_message(&mut self) {
341 self.touch();
342 self.counters.host_messages = self.counters.host_messages.saturating_add(1);
343 }
344
345 fn record_guest_message(&mut self) {
346 self.touch();
347 self.counters.guest_messages = self.counters.guest_messages.saturating_add(1);
348 }
349
350 fn add_exec_output_bytes(&mut self, len: usize) {
351 self.counters.exec_output_bytes =
352 self.counters.exec_output_bytes.saturating_add(len as u64);
353 }
354
355 fn add_fs_bytes(&mut self, len: usize) {
356 self.counters.fs_bytes = self.counters.fs_bytes.saturating_add(len as u64);
357 }
358
359 fn add_tcp_bytes(&mut self, len: usize) {
360 self.counters.tcp_bytes = self.counters.tcp_bytes.saturating_add(len as u64);
361 }
362
363 fn touch(&mut self) {
364 self.activity_seq = self.activity_seq.saturating_add(1);
365 }
366}
367
368impl BulkInputState {
369 fn new(file: File) -> AgentdResult<Self> {
370 Ok(Self {
371 port: AsyncFd::new(file)?,
372 read_buf: vec![0u8; BULK_SERIAL_READ_BUF_SIZE],
373 input: BytesMut::new(),
374 })
375 }
376
377 async fn read_turn(
379 &mut self,
380 read_budget: &mut AgentReadBudget,
381 ) -> AgentdResult<Vec<IncarnatedBulkFrame>> {
382 let buffered = self.drain_turn()?;
383 if !buffered.is_empty() {
384 return Ok(buffered);
385 }
386
387 let mut guard = self.port.readable().await?;
388 match guard
389 .try_io(|inner| read_budget.read_fd(inner.get_ref().as_raw_fd(), &mut self.read_buf))
390 {
391 Ok(Ok(0)) => {
392 return Err(AgentdError::ExecSession(
393 "dedicated bulk port closed".into(),
394 ));
395 }
396 Ok(Ok(read)) => {
397 self.input.extend_from_slice(&self.read_buf[..read]);
398 if self.input.len() > MAX_BULK_INPUT_BUF_SIZE {
399 return Err(AgentdError::ExecSession(
400 "dedicated bulk input exceeded maximum frame buffer".into(),
401 ));
402 }
403 }
404 Ok(Err(error)) if error.kind() == std::io::ErrorKind::Interrupted => {}
405 Ok(Err(error)) => return Err(error.into()),
406 Err(_would_block) => {}
407 }
408 self.drain_turn()
409 }
410
411 fn drain_turn(&mut self) -> AgentdResult<Vec<IncarnatedBulkFrame>> {
412 let mut frames = Vec::new();
413 let mut fs_bytes = 0usize;
414 let mut tcp_bytes = 0usize;
415 while frames.len() < BULK_READER_MAX_RECORDS_PER_TURN
416 && fs_bytes < MAX_BULK_RECORD_PAYLOAD as usize
417 && tcp_bytes < BULK_READER_MAX_BYTES_PER_TURN
418 {
419 let Some(frame) = try_decode_incarnated_bulk_from_bytes(&mut self.input)
420 .map_err(|error| AgentdError::ExecSession(format!("decode agent-bulk: {error}")))?
421 else {
422 break;
423 };
424 match frame.record.kind {
425 BulkKind::Filesystem => {
426 fs_bytes = fs_bytes.saturating_add(frame.record.payload.len());
427 }
428 BulkKind::Tcp => {
429 tcp_bytes = tcp_bytes.saturating_add(frame.record.payload.len());
430 }
431 }
432 frames.push(frame);
433 }
434 Ok(frames)
435 }
436}
437
438impl AgentReadBudget {
439 fn read_fd(&mut self, fd: i32, buf: &mut [u8]) -> std::io::Result<usize> {
440 let result = read_from_fd(fd, buf);
441 self.record_read(result.as_ref().copied().unwrap_or(0));
442 result
443 }
444
445 fn record_read(&mut self, bytes: usize) {
446 self.calls = self.calls.saturating_add(1);
447 self.bytes = self.bytes.saturating_add(bytes);
448 }
449
450 fn exhausted(&self) -> bool {
451 self.bytes >= AGENT_READ_QUANTUM_BYTES || self.calls >= AGENT_READ_QUANTUM_CALLS
452 }
453
454 async fn yield_if_exhausted(&mut self) {
458 if self.exhausted() {
459 tokio::task::yield_now().await;
460 *self = Self::default();
461 }
462 }
463}
464
465pub async fn run(
477 boot_time_ns: u64,
478 init_time_ns: u64,
479 config: &AgentdConfig,
480 port_file: File,
481 boot_console: BootConsoleState,
482 bulk_port: Option<BoundBulkPort>,
483) -> AgentdResult<()> {
484 let process_manager = ProcessManager::get()?;
485 let mut process_manager_failure = process_manager.subscribe_failure()?;
486
487 let port_fd = port_file.as_raw_fd();
490 set_nonblocking(port_fd)?;
491
492 let async_port = AsyncFd::new(port_file)?;
494
495 let mut read_buf = vec![0u8; primary_serial_read_buf_size(bulk_port.is_some())];
499 let mut serial_in_buf = BytesMut::from(boot_console.input.as_slice());
502 let mut serial_out_buf = Vec::new();
503
504 let mut state = AgentState::default();
505 let mut workload = if handoff::is_pid_1() {
506 WorkloadLatch::initialize()
507 } else {
508 WorkloadLatch::unavailable(
509 "PID 1 handoff workloads are not wholly owned by agentd's cgroup",
510 )
511 };
512 if let Some(reason) = workload.unavailable_reason() {
513 eprintln!("checkpoint workload freezer unavailable: {reason}");
514 }
515
516 let (mut session_tx, mut session_rx, bulk_session_rx, bulk_command_rx) =
518 SessionOutputSender::split_channel();
519
520 let mut activity = ActivityTracker::new();
522 let (heartbeat_tx, heartbeat_rx) = watch::channel(heartbeat_snapshot(&state, &activity));
523 let heartbeat_shutdown = Arc::new(AtomicBool::new(false));
530 let heartbeat_control = Arc::new(heartbeat::HeartbeatControl::default());
531 let heartbeat_thread = spawn_heartbeat_thread(
532 heartbeat_rx,
533 Arc::clone(&heartbeat_shutdown),
534 Arc::clone(&heartbeat_control),
535 );
536
537 let ready_time_ns = clock::boottime_ns();
539 let ready_msg = Message::with_payload(
540 MessageType::Ready,
541 0,
542 &Ready {
543 boot_time_ns,
544 init_time_ns,
545 ready_time_ns,
546 agent_version: env!("CARGO_PKG_VERSION").to_string(),
547 bulk_transport: bulk_port
548 .as_ref()
549 .map(|port| BulkTransportReady::dual_port_v1(port.connection_id)),
550 relay_lease: Some(RelayLeaseReady::range_lease_v1()),
551 local_transport: None,
554 workload_transport_barrier_version: Some(WORKLOAD_TRANSPORT_BARRIER_VERSION),
555 },
556 )
557 .map_err(|e| AgentdError::ExecSession(format!("encode ready: {e}")))?;
558 codec::encode_to_buf(&ready_msg, &mut serial_out_buf)
559 .map_err(|e| AgentdError::ExecSession(format!("encode ready frame: {e}")))?;
560 flush_write_buf(&async_port, &mut serial_out_buf).await?;
561
562 let (mut combined_bulk_rx, mut bulk_input, mut bulk_failure_rx, mut bulk_activity_rx) =
563 match bulk_port {
564 Some(port) => {
565 let input = BulkInputState::new(port.file.try_clone()?)?;
566 let (failure_rx, activity_rx) =
567 spawn_bulk_writer_task(port.file, bulk_session_rx, bulk_command_rx);
568 (None, Some(input), Some(failure_rx), Some(activity_rx))
569 }
570 None => {
571 session_tx.disable_bulk_scheduler();
572 (Some(bulk_session_rx), None, None, None)
573 }
574 };
575 let dual_port_active = bulk_input.is_some();
576 let mut pending_bulk_inputs = VecDeque::<PendingBulkInput>::new();
577 let mut bulk_input_bytes_since_snapshot = 0usize;
578 let mut last_bulk_input_snapshot = Instant::now();
579 let input_refunds = state.input_window.clone();
580 let mut credit_deadline = None;
581 let mut last_input_credit = state.input_window.credit()?;
582 let mut read_budget = AgentReadBudget::default();
583
584 'agent: loop {
586 read_budget.yield_if_exhausted().await;
589 if state.resume_output_after_flush {
590 flush_write_buf(&async_port, &mut serial_out_buf).await?;
593 session_tx
594 .resume_bulk_output()
595 .await
596 .map_err(|error| AgentdError::ExecSession(error.into()))?;
597 state.resume_output_after_flush = false;
598 state.output_parked = false;
599 }
600 if state.pending_freeze.as_ref().is_some_and(|message| {
601 message.payload::<WorkloadFreeze>().is_ok_and(|request| {
602 let position = state.input_window.position();
603 position.bulk_bytes >= request.host_input.bulk_bytes
604 && position.bulk_frames >= request.host_input.bulk_frames
605 })
606 }) {
607 let message = state.pending_freeze.take().expect("ready pending freeze");
608 handle_message_with_charge(
609 message,
610 &mut state,
611 &mut activity,
612 &mut session_tx,
613 &mut serial_out_buf,
614 config,
615 &mut workload,
616 &heartbeat_control,
617 None,
618 )
619 .await?;
620 flush_write_buf(&async_port, &mut serial_out_buf).await?;
621 }
622 if !state.output_parked {
623 let credit = state.input_window.credit()?;
624 if input_credit_update_due(&last_input_credit, &credit) {
625 encode_input_credit(credit, &mut serial_out_buf)?;
626 flush_write_buf(&async_port, &mut serial_out_buf).await?;
627 last_input_credit = credit;
628 credit_deadline = None;
629 }
630 }
631 pending_bulk_inputs.retain(|pending| {
634 client_incarnation_for_id(&state, pending.frame.record.id)
635 == Some(pending.frame.incarnation)
636 });
637 let pending_bulk_admission = pending_bulk_inputs.front().map(|pending| {
638 (
639 Arc::clone(&pending.budget),
640 pending.frame.record.payload.len(),
641 )
642 });
643 let has_pending_bulk_admission = pending_bulk_admission.is_some();
644 tokio::select! {
645 _ = input_refunds.refunded(), if !state.output_parked && credit_deadline.is_none() => {
646 if state.input_window.credit()? != last_input_credit {
647 credit_deadline = Some(time::Instant::now() + Duration::from_millis(5));
648 }
649 }
650
651 _ = wait_input_credit_deadline(credit_deadline), if !state.output_parked => {
652 credit_deadline = None;
653 let credit = state.input_window.credit()?;
654 if credit != last_input_credit {
655 encode_input_credit(credit, &mut serial_out_buf)?;
656 flush_write_buf(&async_port, &mut serial_out_buf).await?;
657 last_input_credit = credit;
658 }
659 }
660
661 (id, inherited, result) = std::future::poll_fn(|cx| poll_pending_stdin(&mut state, cx)),
662 if !state.output_parked && !workload.is_frozen() => {
663 if !inherited && let Err(error) = result {
664 encode_stdin_error(id, &AgentdError::Io(error), &mut serial_out_buf)?;
665 flush_write_buf(&async_port, &mut serial_out_buf).await?;
666 }
667 }
668 failure = process_manager_failure.changed() => {
669 let error = match failure {
670 Ok(()) => process_manager_failure
671 .borrow()
672 .clone()
673 .unwrap_or_else(|| "process manager stopped without an error".to_string()),
674 Err(error) => format!("process manager failure channel closed: {error}"),
675 };
676 return Err(AgentdError::ExecSession(error));
677 }
678
679 Some(error) = recv_optional(&mut bulk_failure_rx) => {
680 if state.output_parked {
681 return Err(AgentdError::ExecSession(format!(
682 "dedicated bulk transport failed while frozen: {error}"
683 )));
684 }
685 cancel_all_bulk_correlations(
686 &mut state,
687 &session_tx,
688 &mut serial_out_buf,
689 "dedicated bulk transport failed",
690 )?;
691 let _ = time::timeout(
692 BULK_FAILURE_FLUSH_TIMEOUT,
693 flush_write_buf(&async_port, &mut serial_out_buf),
694 ).await;
695 return Err(AgentdError::ExecSession(format!(
696 "dedicated bulk transport failed: {error}"
697 )));
698 }
699
700 Some(output_activity) = recv_optional(&mut bulk_activity_rx) => {
701 apply_raw_activity(output_activity, &mut activity);
702 publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
703 }
704
705 permit = acquire_bulk_input_permit(pending_bulk_admission), if has_pending_bulk_admission => {
706 let pending = pending_bulk_inputs
707 .pop_front()
708 .expect("guarded pending bulk input exists");
709 let permit = match permit {
710 Ok(permit) => permit,
711 Err(error) => {
712 if validate_bulk_client_incarnation(
716 &state,
717 pending.frame.record.id,
718 pending.frame.incarnation,
719 )? {
720 return Err(error);
721 }
722 continue;
723 }
724 };
725 if !validate_bulk_client_incarnation(
726 &state,
727 pending.frame.record.id,
728 pending.frame.incarnation,
729 )? {
730 continue;
731 }
732 let payload_len = pending.frame.record.payload.len();
733 let bulk_session_tx = session_tx.with_incarnation(Some(pending.frame.incarnation));
734 handle_bulk_record(
735 AdmittedBulkRecord {
736 record: pending.frame.record,
737 _permit: permit,
738 _transport_charge: Some(pending.charge),
739 },
740 &mut state,
741 &mut activity,
742 &mut serial_out_buf,
743 &bulk_session_tx,
744 ).await?;
745 bulk_input_bytes_since_snapshot =
746 bulk_input_bytes_since_snapshot.saturating_add(payload_len);
747 if bulk_input_bytes_since_snapshot >= BULK_ACTIVITY_PUBLISH_BYTES
748 || last_bulk_input_snapshot.elapsed() >= BULK_ACTIVITY_PUBLISH_INTERVAL
749 {
750 publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
751 bulk_input_bytes_since_snapshot = 0;
752 last_bulk_input_snapshot = Instant::now();
753 }
754 if !serial_out_buf.is_empty() {
755 flush_write_buf(&async_port, &mut serial_out_buf).await?;
756 }
757 }
758
759 Some(envelope) = recv_optional(&mut combined_bulk_rx), if !state.output_parked => {
760 if discard_inherited_output(&mut state, &envelope, session_tx.generation()) {
761 continue;
762 }
763 if envelope.incarnation.is_some()
764 && client_incarnation_for_id(&state, envelope.id) != envelope.incarnation
765 {
766 continue;
767 }
768 let SessionOutput::Bulk(output) = envelope.output else {
769 return Err(AgentdError::ExecSession(
770 "non-bulk event entered combined bulk queue".into(),
771 ));
772 };
773 apply_raw_activity(output.activity, &mut activity);
774 if !serial_out_buf.is_empty() {
775 flush_write_buf(&async_port, &mut serial_out_buf).await?;
776 }
777 write_bulk_record_async_fd(&async_port, &output.record).await?;
778 publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
779 }
780
781 turn = async {
782 bulk_input
783 .as_mut()
784 .expect("guarded dedicated bulk input")
785 .read_turn(&mut read_budget)
786 .await
787 }, if bulk_input.is_some() && pending_bulk_inputs.is_empty() => {
788 let frames = match turn {
789 Ok(frames) => frames,
790 Err(error) => {
791 cancel_all_bulk_correlations(
792 &mut state,
793 &session_tx,
794 &mut serial_out_buf,
795 "dedicated bulk transport failed",
796 )?;
797 let _ = time::timeout(
798 BULK_FAILURE_FLUSH_TIMEOUT,
799 flush_write_buf(&async_port, &mut serial_out_buf),
800 ).await;
801 return Err(AgentdError::ExecSession(format!(
802 "dedicated bulk transport failed: {error}"
803 )));
804 }
805 };
806 let mut capacity_deferred = false;
807 if state.output_parked && !frames.is_empty() {
808 return Err(AgentdError::ExecSession("bulk input crossed the frozen transport cut".into()));
809 }
810 for frame in frames {
811 let charge = state.input_window.admit(
812 InputLane::Bulk, bulk_wire_bytes(&frame.record, true),
813 )?;
814 if !validate_bulk_client_incarnation(
815 &state,
816 frame.record.id,
817 frame.incarnation,
818 )? {
819 continue;
822 }
823 let budget = Arc::clone(&state.bulk_input_budget);
824 if capacity_deferred {
825 pending_bulk_inputs.push_back(PendingBulkInput { frame, budget, charge });
826 continue;
827 }
828 let payload_len = frame.record.payload.len();
829 let charged = u32::try_from(payload_len).map_err(|_| {
830 AgentdError::ExecSession("bulk input payload budget overflow".into())
831 })?;
832 let permit = match Arc::clone(&budget).try_acquire_many_owned(charged) {
833 Ok(permit) => permit,
834 Err(_) if !budget.is_closed() => {
835 capacity_deferred = true;
836 pending_bulk_inputs.push_back(PendingBulkInput { frame, budget, charge });
837 continue;
838 }
839 Err(error) => {
840 return Err(AgentdError::ExecSession(format!(
841 "bulk input budget closed unexpectedly: {error}"
842 )));
843 }
844 };
845 let bulk_session_tx = session_tx.with_incarnation(Some(frame.incarnation));
846 handle_bulk_record(
847 AdmittedBulkRecord {
848 record: frame.record,
849 _permit: permit,
850 _transport_charge: Some(charge),
851 },
852 &mut state,
853 &mut activity,
854 &mut serial_out_buf,
855 &bulk_session_tx,
856 ).await?;
857 bulk_input_bytes_since_snapshot =
858 bulk_input_bytes_since_snapshot.saturating_add(payload_len);
859 }
860 if bulk_input_bytes_since_snapshot >= BULK_ACTIVITY_PUBLISH_BYTES
861 || last_bulk_input_snapshot.elapsed() >= BULK_ACTIVITY_PUBLISH_INTERVAL
862 {
863 publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
864 bulk_input_bytes_since_snapshot = 0;
865 last_bulk_input_snapshot = Instant::now();
866 }
867 if !serial_out_buf.is_empty() {
868 flush_write_buf(&async_port, &mut serial_out_buf).await?;
869 }
870 }
871
872 result = async_port.readable() => {
874 let Ok(mut guard) = result else {
875 break;
876 };
877 let mut combined_bulk_records = 0usize;
878 let mut combined_bulk_payload_bytes = 0usize;
879 let mut combined_turn_exhausted = false;
880
881 loop {
882 match guard.try_io(|inner| {
883 read_budget.read_fd(inner.get_ref().as_raw_fd(), &mut read_buf)
884 }) {
885 Ok(Ok(0)) => {
886 if !handoff::is_pid_1() {
888 guard.clear_ready();
889 drop(guard);
890 time::sleep(Duration::from_millis(100)).await;
891 break;
892 }
893 break 'agent;
894 }
895 Ok(Ok(n)) => {
896 serial_in_buf.extend_from_slice(&read_buf[..n]);
897
898 if serial_in_buf.len() > MAX_INPUT_BUF_SIZE {
900 return Err(AgentdError::ExecSession(
901 "serial input buffer exceeded maximum size".into(),
902 ));
903 }
904
905 loop {
910 let input_before = serial_in_buf.len();
911 if let Some(connected) =
912 try_decode_relay_client_connected_from_bytes(&mut serial_in_buf)
913 .map_err(|e| AgentdError::ExecSession(format!(
914 "decode relay client lease: {e}"
915 )))?
916 {
917 if state.output_parked {
918 return Err(AgentdError::ExecSession("relay lease crossed the frozen transport cut".into()));
919 }
920 let _charge = state.input_window.admit(
921 InputLane::Control, input_before - serial_in_buf.len(),
922 )?;
923 establish_relay_client(&mut state, connected)?;
924 continue;
925 }
926 let Some(frame) = codec::try_decode_frame_from_bytes(&mut serial_in_buf)
927 .map_err(|e| AgentdError::ExecSession(format!("decode frame: {e}")))?
928 else {
929 break;
930 };
931 let wire_bytes = input_before - serial_in_buf.len();
932 let DecodedFrame::Control(msg) = frame else {
933 let DecodedFrame::Bulk(record) = frame else {
934 unreachable!();
935 };
936 if dual_port_active {
937 return Err(AgentdError::ExecSession(
938 "raw bulk record arrived on the bound control port"
939 .into(),
940 ));
941 }
942 if state.output_parked {
943 return Err(AgentdError::ExecSession("raw input crossed the frozen transport cut".into()));
944 }
945 let bulk_session_tx = session_tx.with_incarnation(
946 client_incarnation_for_id(&state, record.id),
947 );
948 let payload_len = record.payload.len();
949 let charge = state.input_window.admit(InputLane::Bulk, wire_bytes)?;
950 let budget = Arc::clone(&state.bulk_input_budget);
951 let permit = acquire_bulk_input_permit(Some((
952 budget,
953 payload_len,
954 ))).await?;
955 handle_bulk_record(
956 AdmittedBulkRecord {
957 record,
958 _permit: permit,
959 _transport_charge: Some(charge),
960 },
961 &mut state,
962 &mut activity,
963 &mut serial_out_buf,
964 &bulk_session_tx,
965 ).await?;
966 combined_bulk_records = combined_bulk_records.saturating_add(1);
967 combined_bulk_payload_bytes =
968 combined_bulk_payload_bytes.saturating_add(payload_len);
969 bulk_input_bytes_since_snapshot =
970 bulk_input_bytes_since_snapshot.saturating_add(payload_len);
971 if bulk_input_bytes_since_snapshot >= BULK_ACTIVITY_PUBLISH_BYTES
972 || last_bulk_input_snapshot.elapsed()
973 >= BULK_ACTIVITY_PUBLISH_INTERVAL
974 {
975 publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
976 bulk_input_bytes_since_snapshot = 0;
977 last_bulk_input_snapshot = Instant::now();
978 }
979 if bulk_reader_turn_exhausted(
980 combined_bulk_records,
981 combined_bulk_payload_bytes,
982 ) {
983 combined_turn_exhausted = true;
989 }
990 continue;
991 };
992 let charge = if private_lifecycle_message(&msg) {
993 None
994 } else {
995 if state.output_parked {
996 return Err(AgentdError::ExecSession("ordinary input crossed the frozen transport cut".into()));
997 }
998 let lane = if msg.t.uses_workload_data_credit() {
999 InputLane::Bulk
1000 } else {
1001 InputLane::Control
1002 };
1003 Some(state.input_window.admit(lane, wire_bytes)?)
1004 };
1005 if msg.flags != msg.t.flags() {
1006 let out_before = serial_out_buf.len();
1007 encode_core_error_if_supported(
1008 &msg,
1009 msg.id,
1010 CoreErrorKind::InvalidFlags,
1011 format!(
1012 "invalid flags for {}: got {}, expected {}",
1013 msg.t.as_str(),
1014 msg.flags,
1015 msg.t.flags()
1016 ),
1017 Some(msg.t.as_str().to_string()),
1018 &mut serial_out_buf,
1019 )?;
1020 record_encoded_guest_messages(
1021 &serial_out_buf,
1022 out_before,
1023 &mut activity,
1024 );
1025 publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
1026 continue;
1027 }
1028
1029 if message_refreshes_idle_timer(&msg.t) {
1030 activity.record_host_message();
1031 publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
1032 }
1033
1034 let out_before = serial_out_buf.len();
1035 handle_message_with_charge(
1036 msg,
1037 &mut state,
1038 &mut activity,
1039 &mut session_tx,
1040 &mut serial_out_buf,
1041 config,
1042 &mut workload,
1043 &heartbeat_control,
1044 charge,
1045 ).await?;
1046 record_encoded_guest_messages(
1047 &serial_out_buf,
1048 out_before,
1049 &mut activity,
1050 );
1051 publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
1052 }
1053
1054 if !serial_out_buf.is_empty() {
1056 flush_write_buf(&async_port, &mut serial_out_buf).await?;
1057 }
1058 if combined_turn_exhausted
1062 || state.resume_output_after_flush
1063 || read_budget.exhausted()
1064 {
1065 break;
1066 }
1067 }
1068 Ok(Err(e)) if e.kind() == std::io::ErrorKind::Interrupted => {
1069 if read_budget.exhausted() {
1070 break;
1071 }
1072 continue;
1073 }
1074 Ok(Err(_)) if !handoff::is_pid_1() => {
1075 guard.clear_ready();
1076 drop(guard);
1077 time::sleep(Duration::from_millis(100)).await;
1078 break;
1079 }
1080 Ok(Err(e)) => return Err(e.into()),
1081 Err(_would_block) => break,
1082 }
1083 }
1084 }
1085
1086 Some(envelope) = session_rx.recv(), if !state.output_parked => {
1088 if discard_inherited_output(&mut state, &envelope, session_tx.generation()) {
1089 continue;
1090 }
1091 if envelope.incarnation.is_some()
1092 && client_incarnation_for_id(&state, envelope.id) != envelope.incarnation
1093 {
1094 continue;
1097 }
1098 let id = envelope.id;
1099 match envelope.output {
1100 SessionOutput::Stdout(data) => {
1101 let len = data.len();
1102 let msg = Message::with_payload(MessageType::ExecStdout, id, &ExecStdout { data })
1103 .map_err(|e| AgentdError::ExecSession(format!("encode stdout: {e}")))?;
1104 codec::encode_to_buf(&msg, &mut serial_out_buf)
1105 .map_err(|e| AgentdError::ExecSession(format!("encode stdout frame: {e}")))?;
1106 activity.record_guest_message();
1107 activity.add_exec_output_bytes(len);
1108 }
1109 SessionOutput::Stderr(data) => {
1110 let len = data.len();
1111 let msg = Message::with_payload(MessageType::ExecStderr, id, &ExecStderr { data })
1112 .map_err(|e| AgentdError::ExecSession(format!("encode stderr: {e}")))?;
1113 codec::encode_to_buf(&msg, &mut serial_out_buf)
1114 .map_err(|e| AgentdError::ExecSession(format!("encode stderr frame: {e}")))?;
1115 activity.record_guest_message();
1116 activity.add_exec_output_bytes(len);
1117 }
1118 SessionOutput::Exited(code) => {
1119 let msg = Message::with_payload(MessageType::ExecExited, id, &ExecExited { code })
1120 .map_err(|e| AgentdError::ExecSession(format!("encode exited: {e}")))?;
1121 codec::encode_to_buf(&msg, &mut serial_out_buf)
1122 .map_err(|e| AgentdError::ExecSession(format!("encode exited frame: {e}")))?;
1123 state.sessions.remove(&id);
1124 activity.record_guest_message();
1125 }
1126 SessionOutput::Raw(output) => {
1127 apply_raw_activity(output.activity, &mut activity);
1128 let completion = output.completion;
1129 complete_raw_session(
1130 id,
1131 completion,
1132 &mut state.read_sessions,
1133 &mut state.tcp_sessions,
1134 );
1135 if matches!(completion, Some(RawSessionCompletion::FsWrite))
1136 && let Some(worker) = state.bulk_write_workers.remove(&id)
1137 {
1138 worker.task.abort();
1139 }
1140 if completion.is_some() {
1141 clear_bulk_receive_state(&mut state, id);
1142 }
1143 if !serial_out_buf.is_empty() {
1147 flush_write_buf(&async_port, &mut serial_out_buf).await?;
1148 }
1149 if !output.frame.is_empty() {
1150 write_all_async_fd(&async_port, &output.frame).await?;
1151 }
1152 }
1153 SessionOutput::Bulk(output) => {
1154 apply_raw_activity(output.activity, &mut activity);
1155 if !serial_out_buf.is_empty() {
1156 flush_write_buf(&async_port, &mut serial_out_buf).await?;
1157 }
1158 write_bulk_record_async_fd(&async_port, &output.record).await?;
1159 }
1160 }
1161 publish_heartbeat_snapshot(&heartbeat_tx, &state, &activity);
1162
1163 if !serial_out_buf.is_empty() {
1164 flush_write_buf(&async_port, &mut serial_out_buf).await?;
1165 }
1166 }
1167 }
1168 }
1169
1170 heartbeat_shutdown.store(true, Ordering::Relaxed);
1171 let _ = heartbeat_thread.join();
1172
1173 Ok(())
1174}
1175
1176fn spawn_bulk_writer_task(
1178 writer_file: File,
1179 output_rx: tokio::sync::mpsc::Receiver<SessionOutputEnvelope>,
1180 command_rx: tokio::sync::mpsc::Receiver<BulkOutputCommand>,
1181) -> (
1182 tokio::sync::mpsc::Receiver<String>,
1183 tokio::sync::mpsc::Receiver<RawActivity>,
1184) {
1185 let (failure_tx, failure_rx) = tokio::sync::mpsc::channel(1);
1186 let (activity_tx, activity_rx) = tokio::sync::mpsc::channel(128);
1187
1188 tokio::spawn(async move {
1189 if let Err(error) = bulk_writer_task(writer_file, output_rx, command_rx, activity_tx).await
1190 {
1191 let _ = failure_tx.send(error.to_string()).await;
1192 }
1193 });
1194
1195 (failure_rx, activity_rx)
1196}
1197
1198fn primary_serial_read_buf_size(has_dedicated_bulk_port: bool) -> usize {
1200 if has_dedicated_bulk_port {
1201 CONTROL_SERIAL_READ_BUF_SIZE
1202 } else {
1203 COMBINED_SERIAL_READ_BUF_SIZE
1204 }
1205}
1206
1207fn bulk_reader_turn_exhausted(records: usize, payload_bytes: usize) -> bool {
1209 records >= BULK_READER_MAX_RECORDS_PER_TURN || payload_bytes >= BULK_READER_MAX_BYTES_PER_TURN
1210}
1211
1212async fn recv_optional<T>(receiver: &mut Option<tokio::sync::mpsc::Receiver<T>>) -> Option<T> {
1214 let Some(active) = receiver.as_mut() else {
1215 return std::future::pending().await;
1216 };
1217 match active.recv().await {
1218 Some(value) => Some(value),
1219 None => {
1220 *receiver = None;
1224 std::future::pending().await
1225 }
1226 }
1227}
1228
1229async fn bulk_writer_task(
1231 file: File,
1232 mut output_rx: tokio::sync::mpsc::Receiver<SessionOutputEnvelope>,
1233 mut command_rx: tokio::sync::mpsc::Receiver<BulkOutputCommand>,
1234 activity_tx: tokio::sync::mpsc::Sender<RawActivity>,
1235) -> AgentdResult<()> {
1236 let async_port = AsyncFd::new(file)?;
1237 let mut generation = 0;
1238 let mut flows = HashMap::<(ClientIncarnation, u32), BulkWriteFlow>::new();
1239 let mut active = VecDeque::<(ClientIncarnation, u32)>::new();
1240 let mut retired = HashMap::<ClientIncarnation, Vec<u64>>::new();
1241 let mut retiring_incarnations = HashSet::<ClientIncarnation>::new();
1242 let mut pending_activity = RawActivity::default();
1243 let mut transport = BulkOutputPosition::default();
1244
1245 loop {
1246 let mut cleanups = Vec::new();
1247 tokio::select! {
1248 biased;
1249 Some(command) = command_rx.recv() => {
1250 cleanups.push(apply_bulk_output_command(
1251 command,
1252 &mut generation,
1253 &mut flows,
1254 &mut active,
1255 &mut retired,
1256 &mut retiring_incarnations,
1257 &mut transport,
1258 )?);
1259 }
1260 Some(envelope) = output_rx.recv(), if !transport.parked => {
1261 enqueue_bulk_output(
1262 envelope,
1263 generation,
1264 &mut flows,
1265 &mut active,
1266 &retired,
1267 &retiring_incarnations,
1268 )?;
1269 }
1270 else => break,
1271 }
1272 while let Ok(command) = command_rx.try_recv() {
1273 cleanups.push(apply_bulk_output_command(
1274 command,
1275 &mut generation,
1276 &mut flows,
1277 &mut active,
1278 &mut retired,
1279 &mut retiring_incarnations,
1280 &mut transport,
1281 )?);
1282 }
1283 while let Ok(envelope) = output_rx.try_recv() {
1284 enqueue_bulk_output(
1285 envelope,
1286 generation,
1287 &mut flows,
1288 &mut active,
1289 &retired,
1290 &retiring_incarnations,
1291 )?;
1292 }
1293 complete_bulk_output_cleanups(cleanups, &mut retired, &mut retiring_incarnations);
1294
1295 while !transport.parked && !active.is_empty() {
1296 let round_len = active.len();
1297 let quantum = if round_len == 1 {
1298 BULK_SCHEDULER_MAX_BURST
1299 } else {
1300 BULK_SCHEDULER_QUANTUM
1301 };
1302 for _ in 0..round_len {
1303 if transport.parked || active.is_empty() {
1304 break;
1305 }
1306 let key = active.pop_front().expect("active flow exists");
1307 if let Some(flow) = flows.get_mut(&key) {
1308 flow.deficit = flow
1309 .deficit
1310 .saturating_add(quantum)
1311 .min(BULK_SCHEDULER_MAX_BURST);
1312 }
1313
1314 let mut burst = 0usize;
1315 loop {
1316 let next_len = flows
1317 .get(&key)
1318 .and_then(|flow| flow.queue.front())
1319 .map(bulk_envelope_payload_len)
1320 .transpose()?
1321 .unwrap_or(0);
1322 let can_send = flows.get(&key).is_some_and(|flow| {
1323 next_len != 0
1324 && next_len <= flow.deficit
1325 && burst.saturating_add(next_len) <= BULK_SCHEDULER_MAX_BURST
1326 });
1327 if !can_send {
1328 break;
1329 }
1330
1331 let envelope = {
1332 let flow = flows.get_mut(&key).expect("scheduled flow exists");
1333 let envelope = flow.queue.pop_front().expect("scheduled record exists");
1334 flow.queued_bytes = flow.queued_bytes.saturating_sub(next_len);
1335 flow.deficit = flow.deficit.saturating_sub(next_len);
1336 envelope
1337 };
1338 let incarnation = envelope.incarnation.ok_or_else(|| {
1339 AgentdError::ExecSession(
1340 "dedicated bulk output is missing client incarnation".into(),
1341 )
1342 })?;
1343 let SessionOutput::Bulk(output) = envelope.output else {
1344 unreachable!("bulk scheduler accepted a non-bulk event");
1345 };
1346 let output_activity = output.activity;
1347 write_incarnated_bulk_record_async_fd(&async_port, incarnation, &output.record)
1348 .await?;
1349 transport.wire_bytes = transport
1350 .wire_bytes
1351 .checked_add(bulk_wire_bytes(&output.record, true) as u64)
1352 .ok_or_else(|| {
1353 AgentdError::ExecSession("bulk output counter exhausted".into())
1354 })?;
1355 pending_activity.guest_messages = pending_activity
1356 .guest_messages
1357 .saturating_add(output_activity.guest_messages);
1358 pending_activity.fs_bytes = pending_activity
1359 .fs_bytes
1360 .saturating_add(output_activity.fs_bytes);
1361 pending_activity.tcp_bytes = pending_activity
1362 .tcp_bytes
1363 .saturating_add(output_activity.tcp_bytes);
1364 if pending_activity
1365 .fs_bytes
1366 .saturating_add(pending_activity.tcp_bytes)
1367 >= BULK_ACTIVITY_PUBLISH_BYTES
1368 {
1369 publish_bulk_activity(&activity_tx, &mut pending_activity)?;
1370 }
1371 burst = burst.saturating_add(next_len);
1372 let mut cleanups = Vec::new();
1375 while let Ok(command) = command_rx.try_recv() {
1376 cleanups.push(apply_bulk_output_command(
1377 command,
1378 &mut generation,
1379 &mut flows,
1380 &mut active,
1381 &mut retired,
1382 &mut retiring_incarnations,
1383 &mut transport,
1384 )?);
1385 }
1386 while let Ok(envelope) = output_rx.try_recv() {
1387 enqueue_bulk_output(
1388 envelope,
1389 generation,
1390 &mut flows,
1391 &mut active,
1392 &retired,
1393 &retiring_incarnations,
1394 )?;
1395 }
1396 complete_bulk_output_cleanups(
1397 cleanups,
1398 &mut retired,
1399 &mut retiring_incarnations,
1400 );
1401 if transport.parked {
1402 break;
1403 }
1404 }
1405
1406 if flows.get(&key).is_some_and(|flow| flow.queue.is_empty()) {
1407 flows.remove(&key);
1408 } else if flows.contains_key(&key) && !active.contains(&key) {
1409 active.push_back(key);
1410 }
1411 }
1412
1413 let mut cleanups = Vec::new();
1414 while let Ok(command) = command_rx.try_recv() {
1415 cleanups.push(apply_bulk_output_command(
1416 command,
1417 &mut generation,
1418 &mut flows,
1419 &mut active,
1420 &mut retired,
1421 &mut retiring_incarnations,
1422 &mut transport,
1423 )?);
1424 }
1425 while let Ok(envelope) = output_rx.try_recv() {
1426 enqueue_bulk_output(
1427 envelope,
1428 generation,
1429 &mut flows,
1430 &mut active,
1431 &retired,
1432 &retiring_incarnations,
1433 )?;
1434 }
1435 complete_bulk_output_cleanups(cleanups, &mut retired, &mut retiring_incarnations);
1436 if active.is_empty() && pending_activity.guest_messages != 0 {
1437 publish_bulk_activity(&activity_tx, &mut pending_activity)?;
1438 }
1439 tokio::task::yield_now().await;
1440 }
1441 }
1442
1443 Ok(())
1444}
1445
1446fn publish_bulk_activity(
1448 sender: &tokio::sync::mpsc::Sender<RawActivity>,
1449 pending: &mut RawActivity,
1450) -> AgentdResult<()> {
1451 match sender.try_send(std::mem::take(pending)) {
1452 Ok(()) => Ok(()),
1453 Err(tokio::sync::mpsc::error::TrySendError::Full(activity)) => {
1454 *pending = activity;
1455 Ok(())
1456 }
1457 Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => Err(AgentdError::ExecSession(
1458 "dedicated bulk activity consumer stopped".into(),
1459 )),
1460 }
1461}
1462
1463fn apply_bulk_output_command(
1464 command: BulkOutputCommand,
1465 generation: &mut u64,
1466 flows: &mut HashMap<(ClientIncarnation, u32), BulkWriteFlow>,
1467 active: &mut VecDeque<(ClientIncarnation, u32)>,
1468 retired: &mut HashMap<ClientIncarnation, Vec<u64>>,
1469 retiring_incarnations: &mut HashSet<ClientIncarnation>,
1470 transport: &mut BulkOutputPosition,
1471) -> AgentdResult<BulkOutputCleanup> {
1472 match command {
1473 BulkOutputCommand::Park { completion } => {
1474 transport.parked = true;
1475 Ok(BulkOutputCleanup::Park {
1476 completion,
1477 wire_bytes: transport.wire_bytes,
1478 })
1479 }
1480 BulkOutputCommand::Resume { completion } => {
1481 transport.parked = false;
1482 Ok(BulkOutputCleanup::Flow(completion))
1483 }
1484 BulkOutputCommand::Restore {
1485 generation: next,
1486 completion,
1487 } => {
1488 *generation = next;
1489 flows.clear();
1490 active.clear();
1491 retired.clear();
1492 retiring_incarnations.clear();
1493 Ok(BulkOutputCleanup::Flow(completion))
1494 }
1495 BulkOutputCommand::DropFlow {
1496 incarnation,
1497 id,
1498 completion,
1499 } => {
1500 let key = (incarnation, id);
1501 flows.remove(&key);
1502 active.retain(|active_key| *active_key != key);
1503 retire_bulk_output(retired, incarnation, id)?;
1504 Ok(BulkOutputCleanup::Flow(completion))
1505 }
1506 BulkOutputCommand::DropIncarnation {
1507 incarnation,
1508 completion,
1509 } => {
1510 flows.retain(|(owner, _), _| *owner != incarnation);
1511 active.retain(|(owner, _)| *owner != incarnation);
1512 retiring_incarnations.insert(incarnation);
1513 Ok(BulkOutputCleanup::Incarnation {
1514 incarnation,
1515 completion,
1516 })
1517 }
1518 }
1519}
1520
1521fn complete_bulk_output_cleanups(
1523 cleanups: Vec<BulkOutputCleanup>,
1524 retired: &mut HashMap<ClientIncarnation, Vec<u64>>,
1525 retiring_incarnations: &mut HashSet<ClientIncarnation>,
1526) {
1527 for cleanup in cleanups {
1528 match cleanup {
1529 BulkOutputCleanup::Park {
1530 completion,
1531 wire_bytes,
1532 } => {
1533 let _ = completion.send(wire_bytes);
1534 }
1535 BulkOutputCleanup::Flow(completion) => {
1536 let _ = completion.send(());
1537 }
1538 BulkOutputCleanup::Incarnation {
1539 incarnation,
1540 completion,
1541 } => {
1542 retired.remove(&incarnation);
1543 retiring_incarnations.remove(&incarnation);
1544 let _ = completion.send(());
1545 }
1546 }
1547 }
1548}
1549
1550fn enqueue_bulk_output(
1551 envelope: SessionOutputEnvelope,
1552 generation: u64,
1553 flows: &mut HashMap<(ClientIncarnation, u32), BulkWriteFlow>,
1554 active: &mut VecDeque<(ClientIncarnation, u32)>,
1555 retired: &HashMap<ClientIncarnation, Vec<u64>>,
1556 retiring_incarnations: &HashSet<ClientIncarnation>,
1557) -> AgentdResult<()> {
1558 if envelope.generation != generation {
1559 return Ok(());
1560 }
1561 let id = envelope.id;
1562 let incarnation = envelope.incarnation.ok_or_else(|| {
1563 AgentdError::ExecSession("dedicated bulk output is missing client incarnation".into())
1564 })?;
1565 let payload_len = bulk_envelope_payload_len(&envelope)?;
1566 if id == 0 {
1567 return Err(AgentdError::ExecSession(
1568 "bulk record cannot use correlation ID zero".into(),
1569 ));
1570 }
1571 if payload_len == 0 || payload_len > MAX_BULK_RECORD_PAYLOAD as usize {
1572 return Err(AgentdError::ExecSession(format!(
1573 "bulk output payload {payload_len} exceeds schedulable maximum {MAX_BULK_RECORD_PAYLOAD}"
1574 )));
1575 }
1576
1577 let key = (incarnation, id);
1578 if retiring_incarnations.contains(&incarnation)
1579 || bulk_output_is_retired(retired, incarnation, id)
1580 {
1581 return Ok(());
1584 }
1585 if !flows.contains_key(&key) {
1586 let client_flows = flows
1587 .keys()
1588 .filter(|(owner, _)| *owner == incarnation)
1589 .count();
1590 if client_flows >= BULK_SCHEDULER_MAX_FLOWS_PER_CLIENT {
1591 return Err(AgentdError::ExecSession(format!(
1592 "relay client exceeded active bulk-flow limit for correlation {id}"
1593 )));
1594 }
1595 flows.insert(
1596 key,
1597 BulkWriteFlow {
1598 queue: VecDeque::new(),
1599 queued_bytes: 0,
1600 deficit: 0,
1601 },
1602 );
1603 active.push_back(key);
1604 }
1605
1606 let flow = flows.get_mut(&key).expect("newly inserted flow exists");
1607 let queued_bytes = flow
1608 .queued_bytes
1609 .checked_add(payload_len)
1610 .ok_or_else(|| AgentdError::ExecSession("bulk flow byte budget overflow".into()))?;
1611 if queued_bytes > BULK_SCHEDULER_FLOW_CAPACITY {
1612 return Err(AgentdError::ExecSession(format!(
1613 "bulk flow {id} exceeded queued byte budget"
1614 )));
1615 }
1616 flow.queued_bytes = queued_bytes;
1617 flow.queue.push_back(envelope);
1618 Ok(())
1619}
1620
1621fn bulk_envelope_payload_len(envelope: &SessionOutputEnvelope) -> AgentdResult<usize> {
1622 match &envelope.output {
1623 SessionOutput::Bulk(output) => Ok(output.record.payload.len()),
1624 _ => Err(AgentdError::ExecSession(
1625 "non-bulk event entered dedicated bulk scheduler".into(),
1626 )),
1627 }
1628}
1629
1630pub fn open_and_bind_bulk_port() -> AgentdResult<Option<BoundBulkPort>> {
1632 let cmdline = std::fs::read_to_string("/proc/cmdline")?;
1633 if !cmdline_requests_dual_port(&cmdline) {
1634 return Ok(None);
1635 }
1636
1637 let deadline = Instant::now() + Duration::from_secs(BULK_BINDING_TIMEOUT_SECS);
1638 let port_path = loop {
1639 match serial::find_serial_port(AGENT_BULK_PORT_NAME) {
1640 Ok(path) => break path,
1641 Err(error) if Instant::now() < deadline => {
1642 let _ = error;
1643 std::thread::sleep(Duration::from_millis(10));
1644 }
1645 Err(error) => return Err(error),
1646 }
1647 };
1648 let file = OpenOptions::new().read(true).write(true).open(port_path)?;
1649 let fd = file.as_raw_fd();
1650 set_nonblocking(fd)?;
1651
1652 let connection_id = random_connection_id()?;
1653 write_all_to_fd(fd, &encode_bulk_hello(connection_id), Some(deadline))?;
1654
1655 let mut ack = [0u8; BULK_BINDING_SIZE];
1656 read_exact_from_fd(fd, &mut ack, deadline, "bulk binding acknowledgement")?;
1657 decode_bulk_ack(&ack, connection_id)
1658 .map_err(|error| AgentdError::ExecSession(format!("bind agent-bulk: {error}")))?;
1659
1660 Ok(Some(BoundBulkPort {
1661 file,
1662 connection_id,
1663 }))
1664}
1665
1666pub fn open_serial_port() -> AgentdResult<File> {
1668 let port_path = serial::find_serial_port(AGENT_PORT_NAME)?;
1670
1671 Ok(OpenOptions::new().read(true).write(true).open(&port_path)?)
1674}
1675
1676pub fn receive_bootstrap(port_file: &File) -> AgentdResult<(GuestBootstrap, BootConsoleState)> {
1678 let fd = port_file.as_raw_fd();
1679 set_nonblocking(fd)?;
1680 let deadline = init_ack_deadline();
1681 let mut state = BootConsoleState::default();
1682 let msg = read_boot_message(fd, &mut state, Some(deadline), "guest bootstrap")?;
1683 let bootstrap = decode_bootstrap_message(msg)?;
1684 state.init_failure_ack = bootstrap.init_failure_ack;
1685 Ok((bootstrap, state))
1686}
1687
1688fn decode_bootstrap_message(msg: Message) -> AgentdResult<GuestBootstrap> {
1689 if msg.id != 0 || msg.flags != 0 {
1690 return Err(AgentdError::Config(format!(
1691 "guest bootstrap requires id=0 and flags=0, got id={} flags={}",
1692 msg.id, msg.flags
1693 )));
1694 }
1695 if msg.t != MessageType::Bootstrap {
1696 return Err(AgentdError::Config(format!(
1697 "expected core.bootstrap as first console frame, got {}",
1698 msg.t.as_str()
1699 )));
1700 }
1701 let min_version = MessageType::Bootstrap.min_protocol_version();
1702 if msg.v < min_version {
1703 return Err(AgentdError::Config(format!(
1704 "guest bootstrap requires protocol generation {min_version} or newer, got {}",
1705 msg.v
1706 )));
1707 }
1708
1709 msg.payload::<GuestBootstrap>()
1710 .map_err(|e| AgentdError::Config(format!("decode guest bootstrap payload: {e}")))
1711}
1712
1713pub fn report_init_context(
1715 port_file: &File,
1716 boot_console: &mut BootConsoleState,
1717 default_user: Option<&str>,
1718) -> AgentdResult<()> {
1719 let (uid, gid) = resolve_default_user(default_user)?;
1720 let deadline = init_ack_deadline();
1721 let fd = port_file.as_raw_fd();
1722 set_nonblocking(fd)?;
1723
1724 let msg = Message::with_payload(
1725 MessageType::InitResolved,
1726 0,
1727 &InitResolved {
1728 default_user: ResolvedUser { uid, gid },
1729 },
1730 )
1731 .map_err(|e| AgentdError::ExecSession(format!("encode init context: {e}")))?;
1732
1733 let mut out = Vec::new();
1734 codec::encode_to_buf(&msg, &mut out)
1735 .map_err(|e| AgentdError::ExecSession(format!("encode init context frame: {e}")))?;
1736 write_all_to_fd(fd, &out, Some(deadline))?;
1737 wait_for_init_ack(fd, boot_console, deadline)
1738}
1739
1740pub fn report_init_failure(
1746 port_file: &File,
1747 boot_console: &mut BootConsoleState,
1748 message: &str,
1749 error: &AgentdError,
1750) -> AgentdResult<()> {
1751 let deadline = if boot_console.init_failure_ack {
1752 None
1753 } else {
1754 Some(Instant::now() + std::time::Duration::from_secs(2))
1755 };
1756 let fd = port_file.as_raw_fd();
1757 set_nonblocking(fd)?;
1758 let msg = Message::with_payload(
1759 MessageType::CoreError,
1760 0,
1761 &CoreError {
1762 kind: CoreErrorKind::InitializationFailed,
1763 message: message.to_owned(),
1764 offending_type: None,
1765 init_failure: match error {
1766 AgentdError::UserNotFound(_) => Some(InitFailureReason::UserNotFound),
1767 AgentdError::GroupNotFound(_) => Some(InitFailureReason::GroupNotFound),
1768 _ => None,
1769 },
1770 workload_failure: None,
1771 },
1772 )?;
1773 let mut out = Vec::new();
1774 codec::encode_to_buf(&msg, &mut out)?;
1775 write_all_to_fd(fd, &out, deadline)?;
1776
1777 loop {
1778 let ack = read_boot_message(fd, boot_console, deadline, "init failure ack")?;
1779 if ack.t != MessageType::InitAck || ack.id != 0 || ack.flags != 0 {
1780 return Err(AgentdError::Init(
1781 "expected init failure acknowledgement".into(),
1782 ));
1783 }
1784 if ack.payload::<InitAck>()?.failure {
1785 return Ok(());
1786 }
1787 }
1789}
1790
1791fn ensure_fs_bulk_write_worker(
1796 id: u32,
1797 state: &mut AgentState,
1798 session_tx: &SessionOutputSender,
1799) -> Result<(), String> {
1800 if state.bulk_write_workers.contains_key(&id) {
1801 return Ok(());
1802 }
1803
1804 let session = state
1805 .write_sessions
1806 .remove(&id)
1807 .ok_or_else(|| format!("unknown filesystem write session: {id}"))?;
1808 if !session.is_bulk() {
1809 state.write_sessions.insert(id, session);
1810 return Err("raw bulk record sent to a generation-6 filesystem write".into());
1811 }
1812
1813 let (records, record_rx) = tokio::sync::mpsc::channel(FS_BULK_INPUT_ITEM_CAPACITY);
1814 let (finish, finish_rx) = tokio::sync::mpsc::channel(1);
1817 let output = session_tx.clone();
1818 let task = tokio::spawn(run_fs_bulk_write_worker(
1819 id, session, record_rx, finish_rx, output,
1820 ));
1821 state.bulk_write_workers.insert(
1822 id,
1823 FsBulkWriteWorker {
1824 records,
1825 finish,
1826 task,
1827 },
1828 );
1829 Ok(())
1830}
1831
1832fn enqueue_fs_bulk_record(
1833 id: u32,
1834 record: AdmittedBulkRecord,
1835 state: &mut AgentState,
1836 session_tx: &SessionOutputSender,
1837) -> Result<(), String> {
1838 ensure_fs_bulk_write_worker(id, state, session_tx)?;
1839 let result = state
1840 .bulk_write_workers
1841 .get(&id)
1842 .expect("filesystem bulk worker was just established")
1843 .records
1844 .try_send(record)
1845 .map_err(|error| format!("filesystem bulk input queue is unavailable: {error}"));
1846 if result.is_err()
1847 && let Some(worker) = state.bulk_write_workers.remove(&id)
1848 {
1849 worker.task.abort();
1850 }
1851 result
1852}
1853
1854fn finish_fs_bulk_write(
1855 id: u32,
1856 finish: BulkFinish,
1857 state: &mut AgentState,
1858 session_tx: &SessionOutputSender,
1859) -> Result<(), String> {
1860 ensure_fs_bulk_write_worker(id, state, session_tx)?;
1861 let result = state
1862 .bulk_write_workers
1863 .get(&id)
1864 .expect("filesystem bulk worker was just established")
1865 .finish
1866 .try_send(finish)
1867 .map_err(|error| format!("filesystem bulk finish path is unavailable: {error}"));
1868 if result.is_err()
1869 && let Some(worker) = state.bulk_write_workers.remove(&id)
1870 {
1871 worker.task.abort();
1872 }
1873 result
1874}
1875
1876async fn run_fs_bulk_write_worker(
1877 id: u32,
1878 mut session: FsWriteSession,
1879 mut record_rx: tokio::sync::mpsc::Receiver<AdmittedBulkRecord>,
1880 mut finish_rx: tokio::sync::mpsc::Receiver<BulkFinish>,
1881 output_tx: SessionOutputSender,
1882) {
1883 let mut pending_finish = None;
1884 let mut finish_open = true;
1885 let mut received_offset = 0u64;
1886 let mut pending_activity_bytes = 0usize;
1887 let mut pending_record = None;
1888
1889 loop {
1890 let record = if pending_record.is_some() {
1891 pending_record.take()
1892 } else {
1893 tokio::select! {
1894 biased;
1895 finish = finish_rx.recv(), if pending_finish.is_none() && finish_open => {
1896 match finish {
1897 Some(finish) => pending_finish = Some(finish),
1898 None => finish_open = false,
1899 }
1900 None
1901 }
1902 record = record_rx.recv() => record,
1903 }
1904 };
1905
1906 let Some(record) = record else {
1907 if pending_finish
1908 .as_ref()
1909 .is_some_and(|finish| finish.final_offset <= received_offset)
1910 {
1911 let finish = pending_finish
1912 .take()
1913 .expect("checked pending finish exists");
1914 let mut frame = Vec::new();
1915 if let Err(error) =
1916 fs::handle_fs_bulk_finish(id, finish, &mut session, &mut frame).await
1917 && encode_bulk_fs_failure(id, error, &mut frame).is_err()
1918 {
1919 return;
1920 }
1921 let activity = RawActivity {
1922 guest_messages: 1,
1923 fs_bytes: std::mem::take(&mut pending_activity_bytes),
1924 ..RawActivity::default()
1925 };
1926 let output = crate::session::RawSessionOutput::new(
1927 frame,
1928 activity,
1929 Some(RawSessionCompletion::FsWrite),
1930 );
1931 let _ = output_tx.send(id, SessionOutput::Raw(output)).await;
1932 return;
1933 }
1934 if record_rx.is_closed() {
1935 return;
1936 }
1937 continue;
1938 };
1939
1940 let mut records = Vec::with_capacity(
1944 FS_BULK_WRITE_COALESCE_BYTES.div_ceil(DEFAULT_FILESYSTEM_BULK_RECORD_PAYLOAD as usize),
1945 );
1946 let mut retained_permits = Vec::with_capacity(records.capacity());
1947 let mut payload_len = record.record().payload.len();
1948 let (record, permit) = record.into_parts();
1949 records.push(record);
1950 retained_permits.push(permit);
1951 while payload_len < FS_BULK_WRITE_COALESCE_BYTES {
1952 match record_rx.try_recv() {
1953 Ok(record) => {
1954 let next_len = payload_len.saturating_add(record.record().payload.len());
1955 if next_len > FS_BULK_WRITE_COALESCE_BYTES {
1956 pending_record = Some(record);
1957 break;
1958 }
1959 let (record, permit) = record.into_parts();
1960 payload_len = next_len;
1961 records.push(record);
1962 retained_permits.push(permit);
1963 }
1964 Err(tokio::sync::mpsc::error::TryRecvError::Empty) => break,
1965 Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => break,
1966 }
1967 }
1968
1969 let mut frame = Vec::new();
1970 let last_record = records
1971 .last()
1972 .expect("filesystem bulk batch contains its initial record");
1973 let record_end = last_record
1974 .offset
1975 .saturating_add(last_record.payload.len() as u64);
1976 let mut completed =
1977 match fs::handle_fs_bulk_records(id, &records, &mut session, &mut frame).await {
1978 Ok(completed) => {
1979 if !completed {
1980 received_offset = record_end;
1981 pending_activity_bytes = pending_activity_bytes.saturating_add(payload_len);
1982 }
1983 completed
1984 }
1985 Err(error) => {
1986 if encode_bulk_fs_failure(id, error, &mut frame).is_err() {
1987 return;
1988 }
1989 true
1990 }
1991 };
1992 drop(retained_permits);
1995
1996 if !completed
1997 && pending_finish
1998 .as_ref()
1999 .is_some_and(|finish| finish.final_offset <= received_offset)
2000 {
2001 let finish = pending_finish
2002 .take()
2003 .expect("checked pending finish exists");
2004 completed = match fs::handle_fs_bulk_finish(id, finish, &mut session, &mut frame).await
2005 {
2006 Ok(completed) => completed,
2007 Err(error) => {
2008 if encode_bulk_fs_failure(id, error, &mut frame).is_err() {
2009 return;
2010 }
2011 true
2012 }
2013 };
2014 }
2015
2016 if frame.is_empty() && pending_activity_bytes < FS_ACTIVITY_BATCH_BYTES && !completed {
2020 continue;
2021 }
2022 let activity = RawActivity {
2023 guest_messages: usize::from(!frame.is_empty()),
2024 fs_bytes: std::mem::take(&mut pending_activity_bytes),
2025 ..RawActivity::default()
2026 };
2027 let completion = completed.then_some(RawSessionCompletion::FsWrite);
2028 let output = crate::session::RawSessionOutput::new(frame, activity, completion);
2029 if !output_tx.send(id, SessionOutput::Raw(output)).await || completed {
2030 return;
2031 }
2032 }
2033}
2034
2035async fn handle_bulk_record(
2037 record: AdmittedBulkRecord,
2038 state: &mut AgentState,
2039 activity: &mut ActivityTracker,
2040 out_buf: &mut Vec<u8>,
2041 session_tx: &SessionOutputSender,
2042) -> AgentdResult<()> {
2043 activity.record_host_message();
2044 let id = record.record().id;
2045 let record_end = record
2046 .record()
2047 .offset
2048 .checked_add(record.record().payload.len() as u64)
2049 .ok_or_else(|| AgentdError::ExecSession("bulk record offset overflow".into()))?;
2050 let record_payload_len = record.record().payload.len();
2051 let kind = record.record().kind;
2052 let flow = record.record().flow;
2053 let mut accepted = false;
2054 match kind {
2055 BulkKind::Filesystem => {
2056 if flow != BulkFlow::HostToGuest {
2057 encode_bulk_fs_failure(
2058 id,
2059 "host sent a filesystem record in the guest-to-host flow".into(),
2060 out_buf,
2061 )?;
2062 cancel_bulk_correlation(id, BulkKind::Filesystem, state, session_tx);
2063 return Ok(());
2064 }
2065
2066 if let Err(error) = enqueue_fs_bulk_record(id, record, state, session_tx) {
2067 encode_bulk_fs_failure(id, error, out_buf)?;
2068 cancel_bulk_correlation(id, BulkKind::Filesystem, state, session_tx);
2069 } else {
2070 accepted = true;
2071 }
2072 }
2073 BulkKind::Tcp => {
2074 if flow != BulkFlow::HostToGuest {
2075 encode_bulk_tcp_failure(
2076 id,
2077 "host sent a TCP record in the guest-to-host flow".into(),
2078 out_buf,
2079 )?;
2080 cancel_bulk_correlation(id, BulkKind::Tcp, state, session_tx);
2081 return Ok(());
2082 }
2083 let result = match state.tcp_sessions.get(&id) {
2084 Some(session) => session.write_bulk(record).await,
2085 None => Err(format!("unknown TCP session: {id}")),
2086 };
2087 if let Err(error) = result {
2088 encode_bulk_tcp_failure(id, error, out_buf)?;
2089 cancel_bulk_correlation(id, BulkKind::Tcp, state, session_tx);
2090 } else {
2091 accepted = true;
2092 activity.add_tcp_bytes(record_payload_len);
2093 }
2094 }
2095 }
2096 if accepted {
2097 state.bulk_received_offsets.insert(id, record_end);
2098 if state
2099 .pending_bulk_finishes
2100 .get(&id)
2101 .is_some_and(|finish| finish.final_offset <= record_end)
2102 {
2103 let finish = state
2104 .pending_bulk_finishes
2105 .remove(&id)
2106 .expect("checked pending finish exists");
2107 dispatch_bulk_finish(id, finish, state, out_buf, session_tx).await?;
2108 }
2109 }
2110 Ok(())
2111}
2112
2113async fn dispatch_bulk_finish(
2115 id: u32,
2116 finish: BulkFinish,
2117 state: &mut AgentState,
2118 out_buf: &mut Vec<u8>,
2119 session_tx: &SessionOutputSender,
2120) -> AgentdResult<()> {
2121 match finish.kind {
2122 BulkKind::Filesystem => {
2123 if let Err(error) = finish_fs_bulk_write(id, finish, state, session_tx) {
2124 encode_bulk_fs_failure(id, error, out_buf)?;
2125 cancel_bulk_correlation(id, BulkKind::Filesystem, state, session_tx);
2126 }
2127 }
2128 BulkKind::Tcp => {
2129 let result = match state.tcp_sessions.get(&id) {
2130 Some(session) => session.finish_bulk(finish).await,
2131 None => Err(format!("unknown TCP session: {id}")),
2132 };
2133 if let Err(error) = result {
2134 encode_bulk_tcp_failure(id, error, out_buf)?;
2135 if let Some(session) = state.tcp_sessions.remove(&id) {
2136 session.close();
2137 }
2138 }
2139 }
2140 }
2141 clear_bulk_receive_state(state, id);
2142 Ok(())
2143}
2144
2145fn has_bulk_receive_session(state: &AgentState, id: u32, kind: BulkKind) -> bool {
2147 match kind {
2148 BulkKind::Filesystem => {
2149 state
2150 .write_sessions
2151 .get(&id)
2152 .is_some_and(FsWriteSession::is_bulk)
2153 || state.bulk_write_workers.contains_key(&id)
2154 }
2155 BulkKind::Tcp => state.tcp_sessions.get(&id).is_some_and(TcpSession::is_bulk),
2156 }
2157}
2158
2159fn clear_bulk_receive_state(state: &mut AgentState, id: u32) {
2161 state.bulk_received_offsets.remove(&id);
2162 state.pending_bulk_finishes.remove(&id);
2163}
2164
2165fn cancel_bulk_correlation(
2167 id: u32,
2168 kind: BulkKind,
2169 state: &mut AgentState,
2170 session_tx: &SessionOutputSender,
2171) {
2172 let owner_tx = session_tx.with_incarnation(client_incarnation_for_id(state, id));
2173 match kind {
2177 BulkKind::Filesystem => {
2178 state.write_sessions.remove(&id);
2179 if let Some(worker) = state.bulk_write_workers.remove(&id) {
2180 worker.task.abort();
2181 }
2182 if let Some(session) = state.read_sessions.remove(&id) {
2183 session.abort();
2184 }
2185 }
2186 BulkKind::Tcp => {
2187 if let Some(session) = state.tcp_sessions.remove(&id) {
2188 session.close();
2189 }
2190 }
2191 }
2192 clear_bulk_receive_state(state, id);
2193 if let Err(error) = owner_tx.drop_bulk_flow(id) {
2194 eprintln!("agentd: failed to purge cancelled bulk output id={id}: {error}");
2195 }
2196}
2197
2198fn cancel_all_bulk_correlations(
2199 state: &mut AgentState,
2200 session_tx: &SessionOutputSender,
2201 out_buf: &mut Vec<u8>,
2202 message: &str,
2203) -> AgentdResult<()> {
2204 let mut correlations = HashMap::<u32, BulkKind>::new();
2205 correlations.extend(
2206 state
2207 .read_sessions
2208 .iter()
2209 .filter(|(_, session)| session.is_bulk())
2210 .map(|(id, _)| (*id, BulkKind::Filesystem)),
2211 );
2212 correlations.extend(
2213 state
2214 .write_sessions
2215 .iter()
2216 .filter(|(_, session)| session.is_bulk())
2217 .map(|(id, _)| (*id, BulkKind::Filesystem)),
2218 );
2219 correlations.extend(
2220 state
2221 .bulk_write_workers
2222 .keys()
2223 .map(|id| (*id, BulkKind::Filesystem)),
2224 );
2225 correlations.extend(
2226 state
2227 .tcp_sessions
2228 .iter()
2229 .filter(|(_, session)| session.is_bulk())
2230 .map(|(id, _)| (*id, BulkKind::Tcp)),
2231 );
2232
2233 for (id, kind) in correlations {
2234 encode_bulk_cancel(
2235 id,
2236 kind,
2237 BulkCancelReason::TransportFailure,
2238 message.to_string(),
2239 out_buf,
2240 )?;
2241 cancel_bulk_correlation(id, kind, state, session_tx);
2242 encode_bulk_terminal_failure(id, kind, message.to_string(), out_buf)?;
2243 }
2244 Ok(())
2245}
2246
2247fn clear_bulk_receive_range(state: &mut AgentState, id_start: u32, id_end_exclusive: u32) {
2249 state
2250 .bulk_received_offsets
2251 .retain(|id, _| *id < id_start || *id >= id_end_exclusive);
2252 state
2253 .pending_bulk_finishes
2254 .retain(|id, _| *id < id_start || *id >= id_end_exclusive);
2255}
2256
2257fn client_incarnation_for_id(state: &AgentState, id: u32) -> Option<ClientIncarnation> {
2258 relay_client_slot(id).and_then(|slot| state.client_incarnations.get(&slot).copied())
2259}
2260
2261fn bulk_output_retired_bit(id: u32) -> Option<(usize, u64)> {
2262 let slot = relay_client_slot(id)?;
2263 let (id_start, _) = relay_client_id_range(slot)?;
2264 let local = usize::try_from(id.checked_sub(id_start)?).ok()?;
2265 Some((
2266 local / u64::BITS as usize,
2267 1u64 << (local % u64::BITS as usize),
2268 ))
2269}
2270
2271fn bulk_output_is_retired(
2272 retired: &HashMap<ClientIncarnation, Vec<u64>>,
2273 incarnation: ClientIncarnation,
2274 id: u32,
2275) -> bool {
2276 let Some((word, mask)) = bulk_output_retired_bit(id) else {
2277 return false;
2278 };
2279 retired
2280 .get(&incarnation)
2281 .and_then(|bitmap| bitmap.get(word))
2282 .is_some_and(|bits| bits & mask != 0)
2283}
2284
2285fn retire_bulk_output(
2286 retired: &mut HashMap<ClientIncarnation, Vec<u64>>,
2287 incarnation: ClientIncarnation,
2288 id: u32,
2289) -> AgentdResult<()> {
2290 let (word, mask) = bulk_output_retired_bit(id).ok_or_else(|| {
2291 AgentdError::ExecSession(format!("cannot retire unassigned bulk correlation {id}"))
2292 })?;
2293 let bitmap = retired.entry(incarnation).or_default();
2294 if bitmap.len() <= word {
2295 bitmap.resize(word + 1, 0);
2296 }
2297 bitmap[word] |= mask;
2298 Ok(())
2299}
2300
2301fn validate_bulk_client_incarnation(
2303 state: &AgentState,
2304 id: u32,
2305 claimed: ClientIncarnation,
2306) -> AgentdResult<bool> {
2307 if client_incarnation_for_id(state, id) == Some(claimed) {
2308 return Ok(true);
2309 }
2310 if state
2311 .client_incarnations
2312 .values()
2313 .any(|current| *current == claimed)
2314 {
2315 return Err(AgentdError::ExecSession(format!(
2316 "dedicated bulk record correlation {id} lies outside its client incarnation range"
2317 )));
2318 }
2319 Ok(false)
2320}
2321
2322async fn acquire_bulk_input_permit(
2323 admission: Option<(Arc<Semaphore>, usize)>,
2324) -> AgentdResult<OwnedSemaphorePermit> {
2325 let (budget, payload_len) = admission.expect("bulk admission future is guarded by queue state");
2326 let charged = u32::try_from(payload_len)
2327 .map_err(|_| AgentdError::ExecSession("bulk input payload budget overflow".into()))?;
2328 budget
2329 .acquire_many_owned(charged)
2330 .await
2331 .map_err(|error| AgentdError::ExecSession(format!("bulk input budget closed: {error}")))
2332}
2333
2334fn validate_relay_client_range(id_start: u32, id_end_exclusive: u32) -> AgentdResult<u32> {
2336 canonical_relay_client_range(id_start, id_end_exclusive).ok_or_else(|| {
2337 AgentdError::ExecSession(format!(
2338 "invalid relay client range [{id_start}, {id_end_exclusive})"
2339 ))
2340 })
2341}
2342
2343fn establish_relay_client(
2344 state: &mut AgentState,
2345 connected: RelayClientConnected,
2346) -> AgentdResult<()> {
2347 let slot = validate_relay_client_range(connected.id_start, connected.id_end_exclusive)?;
2348 if connected.incarnation == [0; CLIENT_INCARNATION_SIZE] {
2349 return Err(AgentdError::ExecSession(
2350 "relay client incarnation cannot be zero".into(),
2351 ));
2352 }
2353 if state
2354 .client_incarnations
2355 .get(&slot)
2356 .is_some_and(|current| *current != connected.incarnation)
2357 {
2358 return Err(AgentdError::ExecSession(format!(
2359 "relay client slot {slot} was replaced before acknowledged disconnect"
2360 )));
2361 }
2362 state
2363 .client_incarnations
2364 .insert(slot, connected.incarnation);
2365 Ok(())
2366}
2367
2368fn disconnect_relay_client(
2370 state: &mut AgentState,
2371 disconnected: RelayClientDisconnected,
2372 session_tx: &SessionOutputSender,
2373) -> AgentdResult<(bool, Option<tokio::sync::oneshot::Receiver<()>>)> {
2374 let slot = validate_relay_client_range(disconnected.id_start, disconnected.id_end_exclusive)?;
2375 if let Some(incarnation) = disconnected.incarnation {
2376 if state.client_incarnations.get(&slot) != Some(&incarnation) {
2377 return Ok((false, None));
2378 }
2379 state.client_incarnations.remove(&slot);
2380 } else if state.client_incarnations.contains_key(&slot) {
2381 return Err(AgentdError::ExecSession(
2382 "bound dual-port range cleanup omitted its incarnation".into(),
2383 ));
2384 }
2385 cleanup_relay_client_range(state, disconnected.id_start, disconnected.id_end_exclusive);
2386 let scheduler_cleanup = match disconnected.incarnation {
2387 Some(incarnation) => session_tx
2388 .drop_bulk_incarnation(incarnation)
2389 .map_err(|error| AgentdError::ExecSession(error.into()))?,
2390 None => None,
2391 };
2392 Ok((true, scheduler_cleanup))
2393}
2394
2395fn cleanup_relay_client_range(state: &mut AgentState, id_start: u32, id_end_exclusive: u32) {
2397 state.sessions.retain(|id, session| {
2398 let keep = *id < id_start || *id >= id_end_exclusive;
2399 if !keep {
2400 let _ = session.send_signal(9);
2403 }
2404 keep
2405 });
2406 state.fs.close_owner_range(id_start, id_end_exclusive);
2407 abort_read_sessions_in_owner_range(&mut state.read_sessions, id_start, id_end_exclusive);
2408 state.write_sessions.retain(|_, session| {
2409 let owner_id = session.owner_id();
2410 owner_id < id_start || owner_id >= id_end_exclusive
2411 });
2412 state.bulk_write_workers.retain(|id, worker| {
2413 let keep = *id < id_start || *id >= id_end_exclusive;
2414 if !keep {
2415 worker.task.abort();
2416 }
2417 keep
2418 });
2419 close_tcp_sessions_in_owner_range(&mut state.tcp_sessions, id_start, id_end_exclusive);
2420 clear_bulk_receive_range(state, id_start, id_end_exclusive);
2421}
2422
2423async fn restore_client_state(
2425 state: &mut AgentState,
2426 sender: &mut SessionOutputSender,
2427) -> AgentdResult<()> {
2428 let generation = sender.generation();
2429 state
2430 .detached_sessions
2431 .extend(
2432 std::mem::take(&mut state.sessions)
2433 .into_iter()
2434 .map(|(id, mut session)| {
2435 session.detach_stdin();
2436 ((generation, id), session)
2437 }),
2438 );
2439 state.client_incarnations.clear();
2440 for (_, session) in state.read_sessions.drain() {
2441 session.abort();
2442 }
2443 for (_, worker) in state.bulk_write_workers.drain() {
2444 worker.task.abort();
2445 let _ = worker.task.await;
2446 }
2447 state.write_sessions.clear();
2448 for (_, session) in state.tcp_sessions.drain() {
2449 session.close();
2450 }
2451 state.fs.clear();
2452 state.bulk_received_offsets.clear();
2453 state.pending_bulk_finishes.clear();
2454 time::timeout(Duration::from_secs(5), sender.restore_generation())
2457 .await
2458 .map_err(|_| {
2459 AgentdError::ExecSession(
2460 "restore bulk cleanup timed out; workloads remain frozen".into(),
2461 )
2462 })?
2463 .map_err(|error| AgentdError::ExecSession(error.into()))
2464}
2465
2466fn discard_inherited_output(
2468 state: &mut AgentState,
2469 envelope: &SessionOutputEnvelope,
2470 generation: u64,
2471) -> bool {
2472 if envelope.generation == generation {
2473 return false;
2474 }
2475 if matches!(envelope.output, SessionOutput::Exited(_)) {
2476 state
2477 .detached_sessions
2478 .remove(&(envelope.generation, envelope.id));
2479 }
2480 true
2481}
2482
2483#[allow(clippy::too_many_arguments)]
2486async fn handle_message_with_charge(
2487 msg: Message,
2488 state: &mut AgentState,
2489 activity: &mut ActivityTracker,
2490 root_session_tx: &mut SessionOutputSender,
2491 out_buf: &mut Vec<u8>,
2492 config: &AgentdConfig,
2493 workload: &mut WorkloadLatch,
2494 heartbeat_control: &heartbeat::HeartbeatControl,
2495 mut input_charge: Option<InputCharge>,
2496) -> AgentdResult<()> {
2497 let session_tx = root_session_tx.with_incarnation(client_incarnation_for_id(state, msg.id));
2500 match msg.t {
2501 MessageType::RootDiskPrepare | MessageType::RootDiskGrow => {
2502 let Some(request) = decode_payload_or_core_error::<
2503 microsandbox_protocol::core::RootDiskGrow,
2504 >(&msg, out_buf)?
2505 else {
2506 return Ok(());
2507 };
2508 let apply = msg.t == MessageType::RootDiskGrow;
2509 let result = tokio::task::spawn_blocking(move || {
2510 crate::root_disk::resize(request.size_bytes, apply)
2511 })
2512 .await
2513 .map_err(|e| AgentdError::ExecSession(format!("root grow worker: {e}")))?;
2514 let reply = match result {
2515 Ok(state) => Message::with_payload(MessageType::RootDiskState, msg.id, &state),
2516 Err(message) => Message::with_payload(
2517 MessageType::CoreError,
2518 msg.id,
2519 &CoreError {
2520 kind: CoreErrorKind::CapabilityUnavailable,
2521 message,
2522 offending_type: Some(msg.t.as_str().into()),
2523 init_failure: None,
2524 workload_failure: None,
2525 },
2526 ),
2527 }
2528 .map_err(|e| AgentdError::ExecSession(format!("encode root capacity: {e}")))?;
2529 codec::encode_to_buf(&reply, out_buf).map_err(|e| {
2530 AgentdError::ExecSession(format!("encode root capacity frame: {e}"))
2531 })?;
2532 }
2533
2534 MessageType::Ping => {
2535 let Some(_) = decode_payload_or_core_error::<Ping>(&msg, out_buf)? else {
2536 return Ok(());
2537 };
2538 let reply = Message::with_payload(MessageType::Pong, msg.id, &Pong {})
2539 .map_err(|e| AgentdError::ExecSession(format!("encode pong: {e}")))?;
2540 codec::encode_to_buf(&reply, out_buf)
2541 .map_err(|e| AgentdError::ExecSession(format!("encode pong frame: {e}")))?;
2542 }
2543
2544 MessageType::Touch => {
2545 let Some(_) = decode_payload_or_core_error::<Touch>(&msg, out_buf)? else {
2546 return Ok(());
2547 };
2548 activity.record_host_message();
2549 let reply = Message::with_payload(
2550 MessageType::Touched,
2551 msg.id,
2552 &Touched {
2553 activity_seq: activity.activity_seq,
2554 },
2555 )
2556 .map_err(|e| AgentdError::ExecSession(format!("encode touched: {e}")))?;
2557 codec::encode_to_buf(&reply, out_buf)
2558 .map_err(|e| AgentdError::ExecSession(format!("encode touched frame: {e}")))?;
2559 }
2560
2561 MessageType::WorkloadFreeze => {
2562 let Some(request) = decode_payload_or_core_error::<WorkloadFreeze>(&msg, out_buf)?
2563 else {
2564 return Ok(());
2565 };
2566 let position = state.input_window.position();
2567 if state.pending_freeze.as_ref().is_some_and(|pending| {
2568 pending
2569 .payload::<WorkloadFreeze>()
2570 .is_ok_and(|old| old != request)
2571 }) {
2572 encode_workload_error(
2573 &msg,
2574 &request.attempt_id,
2575 WorkloadLatchError::Conflict("another transport cut is pending".into()),
2576 out_buf,
2577 )?;
2578 return Ok(());
2579 }
2580 if workload.is_frozen() {
2581 if let Err(error) = workload.require_frozen_attempt(&request.attempt_id) {
2582 encode_workload_error(&msg, &request.attempt_id, error, out_buf)?;
2583 return Ok(());
2584 }
2585 if state
2586 .frozen_host_input
2587 .is_some_and(|cut| cut != request.host_input)
2588 {
2589 encode_workload_error(
2590 &msg,
2591 &request.attempt_id,
2592 WorkloadLatchError::Conflict("frozen transport cut cannot change".into()),
2593 out_buf,
2594 )?;
2595 return Ok(());
2596 }
2597 }
2598 if request.host_input.control_bytes != position.control_bytes
2599 || request.host_input.control_frames != position.control_frames
2600 || request.host_input.bulk_bytes < position.bulk_bytes
2601 || request.host_input.bulk_frames < position.bulk_frames
2602 {
2603 encode_workload_error(
2604 &msg,
2605 &request.attempt_id,
2606 WorkloadLatchError::Conflict(
2607 "transport input cut differs from received complete frames".into(),
2608 ),
2609 out_buf,
2610 )?;
2611 return Ok(());
2612 }
2613 state.aborted_transport_attempt = None;
2614 if request.host_input != position {
2615 state.pending_freeze = Some(msg);
2618 return Ok(());
2619 }
2620 state.pending_freeze = None;
2621 let was_frozen = workload.is_frozen();
2622 match workload.freeze(&request.attempt_id) {
2623 Ok(()) => {
2624 if !was_frozen {
2625 state.restored_attempt = None;
2627 }
2628 if let Err(pause_error) = heartbeat_control.pause() {
2629 let rollback = workload.thaw(&request.attempt_id).err();
2630 let message = match rollback {
2631 Some(error) => format!(
2632 "heartbeat checkpoint gate failed: {pause_error}; workload rollback failed: {error}"
2633 ),
2634 None => format!("heartbeat checkpoint gate failed: {pause_error}"),
2635 };
2636 encode_workload_error(
2637 &msg,
2638 &request.attempt_id,
2639 WorkloadLatchError::Io(std::io::Error::other(message)),
2640 out_buf,
2641 )?;
2642 } else {
2643 state.output_parked = true;
2644 state.frozen_host_input = Some(position);
2645 state.frozen_guest_bulk_bytes = root_session_tx
2646 .park_bulk_output()
2647 .await
2648 .map_err(|error| AgentdError::ExecSession(error.into()))?;
2649 let external_mount_tags = request.external_mount_tags.clone();
2653 let external_mounts_synced = if external_mount_tags.is_empty() {
2654 true
2655 } else if state.bulk_write_workers.is_empty() {
2656 let permit = crate::mount_checkpoint::try_start_sync();
2657 if let Some(permit) = permit {
2658 matches!(
2659 time::timeout(
2660 std::time::Duration::from_secs(20),
2661 tokio::task::spawn_blocking(move || {
2662 let _permit = permit;
2665 crate::mount_checkpoint::sync_external_mounts(
2666 external_mount_tags,
2667 )
2668 }),
2669 )
2670 .await,
2671 Ok(Ok(Ok(())))
2672 )
2673 } else {
2674 false
2675 }
2676 } else {
2677 false
2678 };
2679 let reply = Message::with_payload(
2680 MessageType::WorkloadFrozen,
2681 msg.id,
2682 &WorkloadFrozen {
2683 attempt_id: request.attempt_id,
2684 guest_bulk_bytes_target: state.frozen_guest_bulk_bytes,
2685 input_credit: state.input_window.credit()?,
2686 external_mounts_synced,
2687 },
2688 )
2689 .map_err(|error| {
2690 AgentdError::ExecSession(format!(
2691 "encode workload-frozen response: {error}"
2692 ))
2693 })?;
2694 codec::encode_to_buf(&reply, out_buf).map_err(|error| {
2695 AgentdError::ExecSession(format!(
2696 "encode workload-frozen frame: {error}"
2697 ))
2698 })?;
2699 }
2700 }
2701 Err(error) => encode_workload_error(&msg, &request.attempt_id, error, out_buf)?,
2702 }
2703 }
2704
2705 MessageType::WorkloadThaw => {
2706 let Some(request) = decode_payload_or_core_error::<WorkloadThaw>(&msg, out_buf)? else {
2707 return Ok(());
2708 };
2709 if state.pending_freeze.as_ref().is_some_and(|pending| {
2710 pending
2711 .payload::<WorkloadFreeze>()
2712 .is_ok_and(|freeze| freeze.attempt_id != request.attempt_id)
2713 }) {
2714 encode_workload_error(
2715 &msg,
2716 &request.attempt_id,
2717 WorkloadLatchError::Conflict("another transport cut is pending".into()),
2718 out_buf,
2719 )?;
2720 return Ok(());
2721 }
2722 if state.pending_freeze.as_ref().is_some_and(|pending| {
2723 pending
2724 .payload::<WorkloadFreeze>()
2725 .is_ok_and(|freeze| freeze.attempt_id == request.attempt_id)
2726 }) && !workload.is_frozen()
2727 || (!workload.is_frozen()
2728 && state.aborted_transport_attempt.as_deref() == Some(&request.attempt_id))
2729 {
2730 if request.mode != microsandbox_protocol::core::WorkloadThawMode::Continue {
2731 encode_workload_error(
2732 &msg,
2733 &request.attempt_id,
2734 WorkloadLatchError::Conflict(
2735 "restore requires a completed transport cut".into(),
2736 ),
2737 out_buf,
2738 )?;
2739 return Ok(());
2740 }
2741 if let Some(pending) = state.pending_freeze.take() {
2742 encode_workload_error(
2745 &pending,
2746 &request.attempt_id,
2747 WorkloadLatchError::Conflict(
2748 "transport cut aborted by source continuation".into(),
2749 ),
2750 out_buf,
2751 )?;
2752 }
2753 state.aborted_transport_attempt = Some(request.attempt_id.clone());
2754 encode_input_credit(state.input_window.credit()?, out_buf)?;
2755 let reply = Message::with_payload(
2756 MessageType::WorkloadThawed,
2757 msg.id,
2758 &WorkloadThawed {
2759 attempt_id: request.attempt_id,
2760 },
2761 )?;
2762 codec::encode_to_buf(&reply, out_buf)?;
2763 return Ok(());
2764 }
2765 if request.mode == microsandbox_protocol::core::WorkloadThawMode::Restore
2766 && state.restored_attempt.as_deref() != Some(&request.attempt_id)
2767 {
2768 if let Err(error) = workload.require_frozen_attempt(&request.attempt_id) {
2769 encode_workload_error(&msg, &request.attempt_id, error, out_buf)?;
2770 return Ok(());
2771 }
2772 restore_client_state(state, root_session_tx).await?;
2774 state.restored_attempt = Some(request.attempt_id.clone());
2775 }
2776 match workload.thaw(&request.attempt_id) {
2777 Ok(()) => {
2778 heartbeat_control.resume();
2779 encode_input_credit(state.input_window.credit()?, out_buf)?;
2782 state.resume_output_after_flush = true;
2783 state.frozen_host_input = None;
2784 let reply = Message::with_payload(
2785 MessageType::WorkloadThawed,
2786 msg.id,
2787 &WorkloadThawed {
2788 attempt_id: request.attempt_id,
2789 },
2790 )
2791 .map_err(|error| {
2792 AgentdError::ExecSession(format!(
2793 "encode workload-thawed response: {error}"
2794 ))
2795 })?;
2796 codec::encode_to_buf(&reply, out_buf).map_err(|error| {
2797 AgentdError::ExecSession(format!("encode workload-thawed frame: {error}"))
2798 })?;
2799 }
2800 Err(error) => encode_workload_error(&msg, &request.attempt_id, error, out_buf)?,
2801 }
2802 }
2803
2804 MessageType::ExecRequest => {
2805 let Some(mut req) = decode_payload_or_core_error::<ExecRequest>(&msg, out_buf)? else {
2806 return Ok(());
2807 };
2808 if req.cwd.is_none() {
2809 req.cwd = config.default_cwd().map(str::to_string);
2810 }
2811 prepend_scripts_to_path(&mut req);
2812 if workload.is_frozen() {
2813 encode_exec_failed(
2814 msg.id,
2815 ExecFailed {
2816 kind: ExecFailureKind::Other,
2817 errno: None,
2818 errno_name: None,
2819 message: "sandbox workload is frozen for checkpoint activation".into(),
2820 stage: Some("workload_latch".into()),
2821 },
2822 out_buf,
2823 )?;
2824 return Ok(());
2825 }
2826 let workload_placement = match workload.placement() {
2827 Ok(placement) => placement,
2828 Err(error) => {
2829 encode_exec_failed(
2830 msg.id,
2831 ExecFailed {
2832 kind: ExecFailureKind::Other,
2833 errno: None,
2834 errno_name: None,
2835 message: error.to_string(),
2836 stage: Some("workload_cgroup".into()),
2837 },
2838 out_buf,
2839 )?;
2840 return Ok(());
2841 }
2842 };
2843 match ExecSession::spawn(
2844 msg.id,
2845 &req,
2846 session_tx.clone(),
2847 config.user.as_deref(),
2848 config.security_profile,
2849 workload_placement,
2850 ) {
2851 Ok(session) => {
2852 let reply = Message::with_payload(
2853 MessageType::ExecStarted,
2854 msg.id,
2855 &ExecStarted { pid: session.pid() },
2856 )
2857 .map_err(|e| AgentdError::ExecSession(format!("encode started: {e}")))?;
2858 codec::encode_to_buf(&reply, out_buf).map_err(|e| {
2859 AgentdError::ExecSession(format!("encode started frame: {e}"))
2860 })?;
2861 state.sessions.insert(msg.id, session);
2862 }
2863 Err(e) => {
2864 let payload = match &e {
2870 AgentdError::ExecSpawnFailed(p) => p.clone(),
2871 other => ExecFailed {
2872 kind: ExecFailureKind::Other,
2873 errno: None,
2874 errno_name: None,
2875 message: other.to_string(),
2876 stage: None,
2877 },
2878 };
2879 let reply = Message::with_payload(MessageType::ExecFailed, msg.id, &payload)
2880 .map_err(|e| AgentdError::ExecSession(format!("encode failed: {e}")))?;
2881 codec::encode_to_buf(&reply, out_buf).map_err(|e| {
2882 AgentdError::ExecSession(format!("encode failed frame: {e}"))
2883 })?;
2884 eprintln!("failed to spawn exec session {}: {e}", msg.id);
2885 }
2886 }
2887 }
2888
2889 MessageType::ExecStdin => {
2890 let Some(stdin) = decode_payload_or_core_error::<ExecStdin>(&msg, out_buf)? else {
2891 return Ok(());
2892 };
2893 if let Some(session) = state.sessions.get_mut(&msg.id)
2894 && let Err(error) = session.enqueue_stdin(stdin.data, input_charge.take())
2895 {
2896 encode_stdin_error(msg.id, &AgentdError::Io(error), out_buf)?;
2897 }
2898 }
2899
2900 MessageType::ExecResize => {
2901 let Some(resize) = decode_payload_or_core_error::<ExecResize>(&msg, out_buf)? else {
2902 return Ok(());
2903 };
2904 if let Some(session) = state.sessions.get(&msg.id) {
2905 let _ = session.resize(resize.rows, resize.cols);
2906 }
2907 }
2908
2909 MessageType::ExecSignal => {
2910 let Some(signal) = decode_payload_or_core_error::<ExecSignal>(&msg, out_buf)? else {
2911 return Ok(());
2912 };
2913 if let Some(session) = state.sessions.get(&msg.id) {
2914 let _ = session.send_signal(signal.signal);
2915 }
2916 }
2917
2918 MessageType::FsRequest => {
2919 let Some(req) = decode_payload_or_core_error::<FsRequest>(&msg, out_buf)? else {
2920 return Ok(());
2921 };
2922 match fs::handle_fs_request(msg.id, msg.v, req, &mut state.fs, out_buf, &session_tx)
2923 .await
2924 {
2925 Ok(Some(FsStreamSession::Read(rs))) => {
2926 state.read_sessions.insert(msg.id, rs);
2927 }
2928 Ok(Some(FsStreamSession::Write(ws))) => {
2929 state.write_sessions.insert(msg.id, ws);
2930 }
2931 Ok(None) => {}
2932 Err(e) => {
2933 eprintln!("fs request error for {}: {e}", msg.id);
2934 }
2935 }
2936 }
2937
2938 MessageType::FsData => {
2939 let Some(data) = decode_payload_or_core_error::<FsData>(&msg, out_buf)? else {
2940 return Ok(());
2941 };
2942 let len = data.data.len();
2943 if let Some(session) = state.write_sessions.get_mut(&msg.id) {
2944 match fs::handle_fs_data(msg.id, data, session, out_buf).await {
2945 Ok(true) => {
2946 state.write_sessions.remove(&msg.id);
2948 clear_bulk_receive_state(state, msg.id);
2949 }
2950 Ok(false) => {
2951 activity.add_fs_bytes(len);
2952 }
2953 Err(e) => {
2954 eprintln!("fs data error for {}: {e}", msg.id);
2955 state.write_sessions.remove(&msg.id);
2956 clear_bulk_receive_state(state, msg.id);
2957 }
2958 }
2959 } else {
2960 let resp = microsandbox_protocol::fs::FsResponse {
2962 ok: false,
2963 error: Some(format!("unknown write session: {}", msg.id)),
2964 data: None,
2965 };
2966 let reply = Message::with_payload(MessageType::FsResponse, msg.id, &resp)
2967 .map_err(|e| AgentdError::ExecSession(format!("encode fs error: {e}")))?;
2968 codec::encode_to_buf(&reply, out_buf)
2969 .map_err(|e| AgentdError::ExecSession(format!("encode fs error frame: {e}")))?;
2970 }
2971 }
2972
2973 MessageType::BulkCredit => {
2974 let Some(credit) = decode_payload_or_core_error::<BulkCredit>(&msg, out_buf)? else {
2975 return Ok(());
2976 };
2977 match credit.kind {
2978 BulkKind::Filesystem => {
2979 let result = state
2980 .read_sessions
2981 .get(&msg.id)
2982 .ok_or_else(|| format!("unknown filesystem read session: {}", msg.id))
2983 .and_then(|session| session.apply_credit(credit));
2984 if let Err(error) = result {
2985 encode_bulk_fs_failure(msg.id, error, out_buf)?;
2986 cancel_bulk_correlation(msg.id, BulkKind::Filesystem, state, &session_tx);
2987 }
2988 }
2989 BulkKind::Tcp => {
2990 let result = match state.tcp_sessions.get(&msg.id) {
2991 Some(session) => session.apply_credit(credit).await,
2992 None => Ok(()),
2996 };
2997 if let Err(error) = result {
2998 encode_bulk_tcp_failure(msg.id, error, out_buf)?;
2999 cancel_bulk_correlation(msg.id, BulkKind::Tcp, state, &session_tx);
3000 }
3001 }
3002 }
3003 }
3004
3005 MessageType::BulkFinish => {
3006 let Some(finish) = decode_payload_or_core_error::<BulkFinish>(&msg, out_buf)? else {
3007 return Ok(());
3008 };
3009 let received_offset = state
3010 .bulk_received_offsets
3011 .get(&msg.id)
3012 .copied()
3013 .unwrap_or(0);
3014 if has_bulk_receive_session(state, msg.id, finish.kind)
3015 && finish.final_offset > received_offset
3016 {
3017 if state.pending_bulk_finishes.insert(msg.id, finish).is_some() {
3018 encode_bulk_cancel(
3019 msg.id,
3020 finish.kind,
3021 BulkCancelReason::ProtocolState,
3022 "duplicate deferred bulk finish".into(),
3023 out_buf,
3024 )?;
3025 cancel_bulk_correlation(msg.id, finish.kind, state, &session_tx);
3026 }
3027 } else {
3028 dispatch_bulk_finish(msg.id, finish, state, out_buf, &session_tx).await?;
3029 }
3030 }
3031
3032 MessageType::BulkCancel => {
3033 let Some(cancel) = decode_payload_or_core_error::<BulkCancel>(&msg, out_buf)? else {
3034 return Ok(());
3035 };
3036 cancel_bulk_correlation(msg.id, cancel.kind, state, &session_tx);
3037 encode_bulk_terminal_failure(msg.id, cancel.kind, cancel.message, out_buf)?;
3041 }
3042
3043 MessageType::BulkAccepted => {
3044 encode_core_error_if_supported(
3045 &msg,
3046 msg.id,
3047 CoreErrorKind::UnsupportedMessageType,
3048 "host cannot accept a guest-initiated bulk offer".into(),
3049 Some(msg.t.as_str().to_string()),
3050 out_buf,
3051 )?;
3052 }
3053
3054 MessageType::TcpConnect => {
3055 let Some(req) = decode_payload_or_core_error::<TcpConnect>(&msg, out_buf)? else {
3056 return Ok(());
3057 };
3058 if req.bulk.is_some() && msg.v < BULK_PROTOCOL_VERSION {
3059 encode_tcp_failed(
3060 msg.id,
3061 format!("raw bulk offer requires protocol generation {BULK_PROTOCOL_VERSION}"),
3062 out_buf,
3063 )?;
3064 return Ok(());
3065 }
3066 let session = TcpSession::open(msg.id, req, &session_tx);
3069 state.tcp_sessions.insert(msg.id, session);
3070 }
3071
3072 MessageType::TcpData => {
3073 let Some(data) = decode_payload_or_core_error::<TcpData>(&msg, out_buf)? else {
3074 return Ok(());
3075 };
3076 let len = data.data.len();
3077 if let Some(session) = state.tcp_sessions.get(&msg.id) {
3078 if let Err(e) = session
3079 .write_data_charged(data.data, input_charge.take())
3080 .await
3081 {
3082 state.tcp_sessions.remove(&msg.id);
3083 clear_bulk_receive_state(state, msg.id);
3084 encode_tcp_failed(msg.id, e, out_buf)?;
3085 } else {
3086 activity.add_tcp_bytes(len);
3087 }
3088 } else {
3089 encode_tcp_failed(msg.id, format!("unknown TCP session: {}", msg.id), out_buf)?;
3090 }
3091 }
3092
3093 MessageType::TcpEof => {
3094 let Some(_) = decode_payload_or_core_error::<TcpEof>(&msg, out_buf)? else {
3095 return Ok(());
3096 };
3097 if let Some(session) = state.tcp_sessions.get(&msg.id)
3098 && let Err(e) = session.close_write_charged(input_charge.take()).await
3099 {
3100 state.tcp_sessions.remove(&msg.id);
3101 clear_bulk_receive_state(state, msg.id);
3102 encode_tcp_failed(msg.id, e, out_buf)?;
3103 }
3104 }
3105
3106 MessageType::TcpClose => {
3107 let Some(_) = decode_payload_or_core_error::<TcpClose>(&msg, out_buf)? else {
3108 return Ok(());
3109 };
3110 if let Some(session) = state.tcp_sessions.remove(&msg.id) {
3111 session.close();
3112 }
3113 clear_bulk_receive_state(state, msg.id);
3114 }
3115
3116 MessageType::RelayClientDisconnected => {
3117 let Some(disconnected) =
3118 decode_payload_or_core_error::<RelayClientDisconnected>(&msg, out_buf)?
3119 else {
3120 return Ok(());
3121 };
3122 let disconnect_ack =
3124 disconnected
3125 .incarnation
3126 .map(|incarnation| RelayClientDisconnectedAck {
3127 id_start: disconnected.id_start,
3128 id_end_exclusive: disconnected.id_end_exclusive,
3129 incarnation,
3130 });
3131 let (removed, scheduler_cleanup) =
3132 disconnect_relay_client(state, disconnected, &session_tx)?;
3133 if removed && let Some(disconnect_ack) = disconnect_ack {
3134 if let Some(cleanup) = scheduler_cleanup {
3135 let output_tx = session_tx.clone();
3139 tokio::spawn(async move {
3140 if cleanup.await.is_ok() {
3141 let frame =
3142 encode_relay_client_disconnected_ack(disconnect_ack).to_vec();
3143 let _ = output_tx
3144 .send(
3145 0,
3146 SessionOutput::Raw(RawSessionOutput::new(
3147 frame,
3148 RawActivity::default(),
3149 None,
3150 )),
3151 )
3152 .await;
3153 }
3154 });
3155 } else {
3156 out_buf
3159 .extend_from_slice(&encode_relay_client_disconnected_ack(disconnect_ack));
3160 }
3161 }
3162 }
3163
3164 MessageType::ClockSync => {
3165 let Some(sync) = decode_payload_or_core_error::<ClockSync>(&msg, out_buf)? else {
3166 return Ok(());
3167 };
3168 if let Err(e) = clock::sync_realtime_unix_nanos(sync.unix_time_nanos) {
3169 eprintln!("clock: failed to sync realtime clock: {e}");
3170 }
3171 }
3172
3173 MessageType::Shutdown => {
3174 for session in state
3178 .sessions
3179 .drain()
3180 .map(|(_, session)| session)
3181 .chain(state.detached_sessions.drain().map(|(_, session)| session))
3182 {
3183 let _ = session.send_signal(15); }
3185 state.write_sessions.clear();
3186 for (_, worker) in state.bulk_write_workers.drain() {
3187 worker.task.abort();
3188 }
3189 for (_, session) in state.tcp_sessions.drain() {
3190 session.close();
3191 }
3192 state.fs.clear();
3193
3194 if let Err(error) = request_guest_poweroff() {
3198 eprintln!("agentd: graceful poweroff request failed: {error}");
3199 }
3200 }
3201
3202 _ => {
3203 }
3205 }
3206
3207 Ok(())
3208}
3209
3210fn message_refreshes_idle_timer(t: &MessageType) -> bool {
3221 !matches!(
3222 t,
3223 MessageType::ClockSync
3224 | MessageType::Ping
3225 | MessageType::Touch
3226 | MessageType::WorkloadFreeze
3227 | MessageType::WorkloadThaw
3228 )
3229}
3230
3231fn guest_message_refreshes_idle_timer(t: &MessageType) -> bool {
3239 !matches!(
3240 t,
3241 MessageType::Pong
3242 | MessageType::Touched
3243 | MessageType::WorkloadFrozen
3244 | MessageType::WorkloadThawed
3245 | MessageType::WorkloadTransportCredit
3246 | MessageType::CoreError
3247 )
3248}
3249
3250fn spawn_heartbeat_thread(
3260 snapshot_rx: watch::Receiver<HeartbeatSnapshot>,
3261 shutdown: Arc<AtomicBool>,
3262 control: Arc<heartbeat::HeartbeatControl>,
3263) -> std::thread::JoinHandle<()> {
3264 std::thread::Builder::new()
3265 .name("agentd-heartbeat".to_string())
3266 .spawn(move || {
3267 let mut heartbeat_seq = 0u64;
3268 let mut last_activity_seq = snapshot_rx.borrow().activity_seq;
3269 let mut last_activity = Utc::now();
3270
3271 let interval = Duration::from_secs(HEARTBEAT_INTERVAL_SECS);
3272 let step = Duration::from_millis(100);
3273
3274 while !shutdown.load(Ordering::Relaxed) {
3275 let mut slept = Duration::ZERO;
3276 while slept < interval {
3277 if shutdown.load(Ordering::Relaxed) {
3278 return;
3279 }
3280 std::thread::sleep(step);
3281 slept += step;
3282 }
3283
3284 if !heartbeat::heartbeat_dir_exists() {
3285 continue;
3286 }
3287
3288 heartbeat_seq = heartbeat_seq.saturating_add(1);
3289 let snapshot = snapshot_rx.borrow().clone();
3290 let timestamp = Utc::now();
3291 if snapshot.activity_seq != last_activity_seq {
3292 last_activity_seq = snapshot.activity_seq;
3293 last_activity = timestamp;
3294 }
3295 let heartbeat = Heartbeat {
3296 heartbeat_seq,
3297 activity_seq: snapshot.activity_seq,
3298 timestamp,
3299 last_activity,
3300 active_exec_sessions: snapshot.active_exec_sessions,
3301 active_fs_streams: snapshot.active_fs_streams,
3302 active_tcp_streams: snapshot.active_tcp_streams,
3303 activity_counters: snapshot.counters,
3304 };
3305 let Some(_write_guard) = control.begin_write() else {
3306 continue;
3307 };
3308 let _ = heartbeat::write_heartbeat(&heartbeat);
3309 }
3310 })
3311 .expect("failed to spawn agentd heartbeat thread")
3312}
3313
3314fn heartbeat_snapshot(state: &AgentState, activity: &ActivityTracker) -> HeartbeatSnapshot {
3315 HeartbeatSnapshot {
3316 activity_seq: activity.activity_seq,
3317 active_exec_sessions: (state.sessions.len() + state.detached_sessions.len()) as u32,
3318 active_fs_streams: state
3319 .read_sessions
3320 .len()
3321 .saturating_add(state.write_sessions.len())
3322 .saturating_add(state.bulk_write_workers.len()) as u32,
3323 active_tcp_streams: state.tcp_sessions.len() as u32,
3324 counters: activity.counters,
3325 }
3326}
3327
3328fn initial_input_credit() -> WorkloadTransportCredit {
3329 WorkloadTransportCredit {
3330 control_bytes: WORKLOAD_TRANSPORT_CONTROL_BYTES,
3331 control_frames: WORKLOAD_TRANSPORT_CONTROL_FRAMES,
3332 bulk_bytes: WORKLOAD_TRANSPORT_BULK_BYTES,
3333 bulk_frames: WORKLOAD_TRANSPORT_BULK_FRAMES,
3334 }
3335}
3336
3337fn bulk_wire_bytes(record: &BulkRecord, dedicated: bool) -> usize {
3338 4 + FRAME_HEADER_SIZE
3341 + BULK_HEADER_SIZE
3342 + record.payload.len()
3343 + if dedicated {
3344 CLIENT_INCARNATION_SIZE
3345 } else {
3346 0
3347 }
3348}
3349
3350fn private_lifecycle_message(message: &Message) -> bool {
3351 message.id == u32::MAX
3352 && matches!(
3353 message.t,
3354 MessageType::WorkloadFreeze | MessageType::WorkloadThaw
3355 )
3356 && message.flags == message.t.flags()
3357}
3358
3359fn input_credit_update_due(
3360 previous: &WorkloadTransportCredit,
3361 current: &WorkloadTransportCredit,
3362) -> bool {
3363 current.control_bytes.saturating_sub(previous.control_bytes) >= 64 * 1024
3364 || current.bulk_bytes.saturating_sub(previous.bulk_bytes) >= 64 * 1024
3365 || current
3366 .control_frames
3367 .saturating_sub(previous.control_frames)
3368 >= 16
3369 || current.bulk_frames.saturating_sub(previous.bulk_frames) >= 16
3370}
3371
3372async fn wait_input_credit_deadline(deadline: Option<time::Instant>) {
3373 match deadline {
3374 Some(deadline) => time::sleep_until(deadline).await,
3375 None => std::future::pending().await,
3376 }
3377}
3378
3379fn encode_input_credit(credit: WorkloadTransportCredit, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3380 let message = Message::with_payload(MessageType::WorkloadTransportCredit, u32::MAX, &credit)?;
3383 codec::encode_to_buf(&message, out_buf)?;
3384 Ok(())
3385}
3386
3387fn poll_pending_stdin(
3388 state: &mut AgentState,
3389 cx: &mut Context<'_>,
3390) -> Poll<(u32, bool, std::io::Result<()>)> {
3391 let total = state.sessions.len() + state.detached_sessions.len();
3394 let start = state.stdin_poll_offset.checked_rem(total).unwrap_or(0);
3395 for pass in 0..2 {
3396 let sessions = state
3397 .sessions
3398 .iter_mut()
3399 .map(|(id, session)| (*id, false, session))
3400 .chain(
3401 state
3402 .detached_sessions
3403 .iter_mut()
3404 .map(|((_, id), session)| (*id, true, session)),
3405 );
3406 for (index, (id, inherited, session)) in sessions.enumerate() {
3407 if (pass == 0 && index < start) || (pass == 1 && index >= start) {
3408 continue;
3409 }
3410 if session.has_pending_stdin()
3411 && let Poll::Ready(result) = session.poll_pending_stdin(cx)
3412 {
3413 state.stdin_poll_offset = index + 1;
3414 return Poll::Ready((id, inherited, result));
3415 }
3416 }
3417 }
3418 Poll::Pending
3419}
3420
3421fn encode_stdin_error(id: u32, error: &AgentdError, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3422 let message =
3423 Message::with_payload(MessageType::ExecStdinError, id, &stdin_error_payload(error))?;
3424 codec::encode_to_buf(&message, out_buf)?;
3425 Ok(())
3426}
3427
3428fn publish_heartbeat_snapshot(
3429 heartbeat_tx: &watch::Sender<HeartbeatSnapshot>,
3430 state: &AgentState,
3431 activity: &ActivityTracker,
3432) {
3433 let _ = heartbeat_tx.send(heartbeat_snapshot(state, activity));
3434}
3435
3436fn record_encoded_guest_messages(out_buf: &[u8], start: usize, activity: &mut ActivityTracker) {
3437 let mut offset = start;
3438 while offset + 4 <= out_buf.len() {
3439 let frame_len = u32::from_be_bytes([
3440 out_buf[offset],
3441 out_buf[offset + 1],
3442 out_buf[offset + 2],
3443 out_buf[offset + 3],
3444 ]) as usize;
3445 let total = 4usize.saturating_add(frame_len);
3446 if offset.saturating_add(total) > out_buf.len() {
3447 break;
3448 }
3449
3450 if encoded_guest_message_refreshes_idle_timer(out_buf, offset, frame_len) {
3451 activity.record_guest_message();
3452 }
3453 offset += total;
3454 }
3455}
3456
3457fn encoded_guest_message_refreshes_idle_timer(
3458 out_buf: &[u8],
3459 offset: usize,
3460 frame_len: usize,
3461) -> bool {
3462 let frame_end = offset.saturating_add(4).saturating_add(frame_len);
3463 if frame_end <= out_buf.len() {
3464 let mut frame = BytesMut::from(&out_buf[offset..frame_end]);
3465 if try_decode_relay_client_disconnected_ack_from_bytes(&mut frame)
3466 .is_ok_and(|ack| ack.is_some())
3467 {
3468 return false;
3469 }
3470 }
3471 if frame_len < microsandbox_protocol::message::FRAME_HEADER_SIZE {
3472 return true;
3473 }
3474
3475 let id_start = offset + 4;
3476 let flags_index = id_start + 4;
3477 let body_start = flags_index + 1;
3478 let body_end = offset + 4 + frame_len;
3479 if body_end > out_buf.len() || body_start > body_end {
3480 return true;
3481 }
3482
3483 let id = u32::from_be_bytes([
3484 out_buf[id_start],
3485 out_buf[id_start + 1],
3486 out_buf[id_start + 2],
3487 out_buf[id_start + 3],
3488 ]);
3489 let frame = codec::RawFrame {
3490 id,
3491 flags: out_buf[flags_index],
3492 body: out_buf[body_start..body_end].to_vec(),
3493 };
3494
3495 codec::raw_frame_to_message(frame)
3496 .map(|msg| guest_message_refreshes_idle_timer(&msg.t))
3497 .unwrap_or(true)
3498}
3499
3500fn apply_raw_activity(raw: RawActivity, activity: &mut ActivityTracker) {
3501 if raw.guest_messages != 0 {
3502 activity.activity_seq = activity
3503 .activity_seq
3504 .saturating_add(raw.guest_messages as u64);
3505 activity.counters.guest_messages = activity
3506 .counters
3507 .guest_messages
3508 .saturating_add(raw.guest_messages as u64);
3509 }
3510 if raw.fs_bytes > 0 {
3511 activity.add_fs_bytes(raw.fs_bytes);
3512 }
3513 if raw.tcp_bytes > 0 {
3514 activity.add_tcp_bytes(raw.tcp_bytes);
3515 }
3516}
3517
3518fn complete_raw_session(
3519 id: u32,
3520 completion: Option<RawSessionCompletion>,
3521 read_sessions: &mut HashMap<u32, FsReadSession>,
3522 tcp_sessions: &mut HashMap<u32, TcpSession>,
3523) {
3524 match completion {
3525 Some(RawSessionCompletion::FsRead) => {
3526 read_sessions.remove(&id);
3527 }
3528 Some(RawSessionCompletion::FsWrite) => {}
3529 Some(RawSessionCompletion::Tcp) => {
3530 tcp_sessions.remove(&id);
3531 }
3532 None => {}
3533 }
3534}
3535
3536fn abort_read_sessions_in_owner_range(
3537 read_sessions: &mut HashMap<u32, FsReadSession>,
3538 id_start: u32,
3539 id_end_exclusive: u32,
3540) {
3541 let mut retained = HashMap::new();
3542 for (id, session) in read_sessions.drain() {
3543 let owner_id = session.owner_id();
3544 if owner_id >= id_start && owner_id < id_end_exclusive {
3545 session.abort();
3546 } else {
3547 retained.insert(id, session);
3548 }
3549 }
3550 *read_sessions = retained;
3551}
3552
3553fn close_tcp_sessions_in_owner_range(
3554 tcp_sessions: &mut HashMap<u32, TcpSession>,
3555 id_start: u32,
3556 id_end_exclusive: u32,
3557) {
3558 let mut retained = HashMap::new();
3559 for (id, session) in tcp_sessions.drain() {
3560 let owner_id = session.owner_id();
3561 if owner_id >= id_start && owner_id < id_end_exclusive {
3562 session.close();
3563 } else {
3564 retained.insert(id, session);
3565 }
3566 }
3567 *tcp_sessions = retained;
3568}
3569
3570fn encode_tcp_failed(id: u32, error: String, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3571 let reply = Message::with_payload(MessageType::TcpFailed, id, &TcpFailed { error })
3572 .map_err(|e| AgentdError::ExecSession(format!("encode tcp failed: {e}")))?;
3573 codec::encode_to_buf(&reply, out_buf)
3574 .map_err(|e| AgentdError::ExecSession(format!("encode tcp failed frame: {e}")))?;
3575 Ok(())
3576}
3577
3578fn encode_bulk_fs_failure(id: u32, error: String, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3579 encode_bulk_cancel(
3580 id,
3581 BulkKind::Filesystem,
3582 BulkCancelReason::ProtocolState,
3583 error.clone(),
3584 out_buf,
3585 )?;
3586 let response = Message::with_payload(
3587 MessageType::FsResponse,
3588 id,
3589 &FsResponse {
3590 ok: false,
3591 error: Some(error),
3592 data: None,
3593 },
3594 )
3595 .map_err(|error| AgentdError::ExecSession(format!("encode fs failure: {error}")))?;
3596 codec::encode_to_buf(&response, out_buf)
3597 .map_err(|error| AgentdError::ExecSession(format!("encode fs failure frame: {error}")))
3598}
3599
3600fn encode_bulk_tcp_failure(id: u32, error: String, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3601 encode_bulk_cancel(
3602 id,
3603 BulkKind::Tcp,
3604 BulkCancelReason::ProtocolState,
3605 error.clone(),
3606 out_buf,
3607 )?;
3608 encode_tcp_failed(id, error, out_buf)
3609}
3610
3611fn encode_bulk_terminal_failure(
3613 id: u32,
3614 kind: BulkKind,
3615 error: String,
3616 out_buf: &mut Vec<u8>,
3617) -> AgentdResult<()> {
3618 match kind {
3619 BulkKind::Filesystem => {
3620 let response = Message::with_payload(
3621 MessageType::FsResponse,
3622 id,
3623 &FsResponse {
3624 ok: false,
3625 error: Some(error),
3626 data: None,
3627 },
3628 )
3629 .map_err(|error| {
3630 AgentdError::ExecSession(format!("encode filesystem terminal failure: {error}"))
3631 })?;
3632 codec::encode_to_buf(&response, out_buf).map_err(|error| {
3633 AgentdError::ExecSession(format!(
3634 "encode filesystem terminal failure frame: {error}"
3635 ))
3636 })
3637 }
3638 BulkKind::Tcp => encode_tcp_failed(id, error, out_buf),
3639 }
3640}
3641
3642fn encode_bulk_cancel(
3643 id: u32,
3644 kind: BulkKind,
3645 reason: BulkCancelReason,
3646 message: String,
3647 out_buf: &mut Vec<u8>,
3648) -> AgentdResult<()> {
3649 let cancel = Message::with_payload(
3650 MessageType::BulkCancel,
3651 id,
3652 &BulkCancel {
3653 kind,
3654 reason,
3655 message,
3656 },
3657 )
3658 .map_err(|error| AgentdError::ExecSession(format!("encode bulk cancel: {error}")))?;
3659 codec::encode_to_buf(&cancel, out_buf)
3660 .map_err(|error| AgentdError::ExecSession(format!("encode bulk cancel frame: {error}")))
3661}
3662
3663fn encode_core_error_if_supported(
3664 source: &Message,
3665 id: u32,
3666 kind: CoreErrorKind,
3667 message: String,
3668 offending_type: Option<String>,
3669 out_buf: &mut Vec<u8>,
3670) -> AgentdResult<()> {
3671 if !MessageType::CoreError.is_available_at(source.v) {
3672 return Err(AgentdError::ExecSession(format!(
3673 "cannot send core.error to protocol generation {}",
3674 source.v
3675 )));
3676 }
3677
3678 encode_core_error(id, kind, message, offending_type, out_buf)
3679}
3680
3681fn encode_core_error(
3682 id: u32,
3683 kind: CoreErrorKind,
3684 message: String,
3685 offending_type: Option<String>,
3686 out_buf: &mut Vec<u8>,
3687) -> AgentdResult<()> {
3688 let reply = Message::with_payload(
3689 MessageType::CoreError,
3690 id,
3691 &CoreError {
3692 kind,
3693 message,
3694 offending_type,
3695 init_failure: None,
3696 workload_failure: None,
3697 },
3698 )
3699 .map_err(|e| AgentdError::ExecSession(format!("encode core error: {e}")))?;
3700 codec::encode_to_buf(&reply, out_buf)
3701 .map_err(|e| AgentdError::ExecSession(format!("encode core error frame: {e}")))?;
3702 Ok(())
3703}
3704
3705fn encode_workload_error(
3706 source: &Message,
3707 attempt_id: &str,
3708 error: WorkloadLatchError,
3709 out_buf: &mut Vec<u8>,
3710) -> AgentdResult<()> {
3711 let kind = match &error {
3712 WorkloadLatchError::Unavailable(_) | WorkloadLatchError::Io(_) => {
3713 CoreErrorKind::CapabilityUnavailable
3714 }
3715 WorkloadLatchError::InvalidAttempt(_) => CoreErrorKind::InvalidPayload,
3716 WorkloadLatchError::Conflict(_) => CoreErrorKind::InvalidSession,
3717 };
3718 let disposition = match &error {
3721 WorkloadLatchError::Unavailable(_) => WorkloadFailureDisposition::Unavailable,
3722 _ => WorkloadFailureDisposition::RecoveryRequired,
3723 };
3724 if !MessageType::CoreError.is_available_at(source.v) {
3725 return Err(AgentdError::ExecSession(
3726 "peer cannot receive workload errors".into(),
3727 ));
3728 }
3729 let reply = Message::with_payload(
3730 MessageType::CoreError,
3731 source.id,
3732 &CoreError {
3733 kind,
3734 message: error.to_string(),
3735 offending_type: Some(source.t.as_str().to_string()),
3736 init_failure: None,
3737 workload_failure: Some(WorkloadFailure {
3738 attempt_id: attempt_id.to_string(),
3739 disposition,
3740 }),
3741 },
3742 )
3743 .map_err(|error| AgentdError::ExecSession(format!("encode workload error: {error}")))?;
3744 codec::encode_to_buf(&reply, out_buf)
3745 .map_err(|error| AgentdError::ExecSession(format!("encode workload error frame: {error}")))
3746}
3747
3748fn encode_exec_failed(id: u32, payload: ExecFailed, out_buf: &mut Vec<u8>) -> AgentdResult<()> {
3749 let reply = Message::with_payload(MessageType::ExecFailed, id, &payload)
3750 .map_err(|error| AgentdError::ExecSession(format!("encode exec failure: {error}")))?;
3751 codec::encode_to_buf(&reply, out_buf)
3752 .map_err(|error| AgentdError::ExecSession(format!("encode exec failure frame: {error}")))?;
3753 Ok(())
3754}
3755
3756fn decode_payload_or_core_error<T>(msg: &Message, out_buf: &mut Vec<u8>) -> AgentdResult<Option<T>>
3757where
3758 T: serde::de::DeserializeOwned,
3759{
3760 match msg.payload::<T>() {
3761 Ok(payload) => Ok(Some(payload)),
3762 Err(error) => {
3763 encode_core_error_if_supported(
3764 msg,
3765 msg.id,
3766 CoreErrorKind::InvalidPayload,
3767 format!("decode payload for {}: {error}", msg.t.as_str()),
3768 Some(msg.t.as_str().to_string()),
3769 out_buf,
3770 )?;
3771 Ok(None)
3772 }
3773 }
3774}
3775
3776fn stdin_error_payload(err: &AgentdError) -> ExecStdinError {
3778 let io_err = match err {
3779 AgentdError::Io(e) => Some(e),
3780 _ => None,
3781 };
3782 let errno = io_err.and_then(|e| e.raw_os_error());
3783 ExecStdinError {
3784 errno,
3785 errno_name: errno.and_then(errno_name),
3786 message: err.to_string(),
3787 }
3788}
3789
3790fn errno_name(code: i32) -> Option<String> {
3793 let name = match code {
3794 libc::EPIPE => "EPIPE",
3795 libc::EBADF => "EBADF",
3796 libc::EINVAL => "EINVAL",
3797 libc::EIO => "EIO",
3798 libc::ENOSPC => "ENOSPC",
3799 libc::EFBIG => "EFBIG",
3800 _ => return None,
3801 };
3802 Some(name.to_string())
3803}
3804
3805fn prepend_scripts_to_path(req: &mut microsandbox_protocol::exec::ExecRequest) {
3806 if let Some(entry) = req.env.iter_mut().find(|e| e.starts_with("PATH=")) {
3808 let existing = &entry["PATH=".len()..];
3809 *entry = format!("PATH={}", scripts_path(Some(existing)));
3810 } else {
3811 let inherited = env::var("PATH").ok();
3812 req.env
3813 .push(format!("PATH={}", scripts_path(inherited.as_deref())));
3814 }
3815}
3816
3817fn set_nonblocking(fd: i32) -> AgentdResult<()> {
3819 let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
3820 if flags < 0 {
3821 return Err(std::io::Error::last_os_error().into());
3822 }
3823 let ret = unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) };
3824 if ret < 0 {
3825 return Err(std::io::Error::last_os_error().into());
3826 }
3827 Ok(())
3828}
3829
3830fn cmdline_requests_dual_port(cmdline: &str) -> bool {
3831 cmdline
3832 .split_ascii_whitespace()
3833 .any(|argument| argument == AGENT_TRANSPORT_DUAL_PORT_CMDLINE)
3834}
3835
3836fn random_connection_id() -> AgentdResult<[u8; 16]> {
3837 let mut id = [0u8; 16];
3838 let mut filled = 0;
3839 while filled < id.len() {
3840 let result =
3841 unsafe { libc::getrandom(id[filled..].as_mut_ptr().cast(), id.len() - filled, 0) };
3842 if result > 0 {
3843 filled += result as usize;
3844 continue;
3845 }
3846 let error = std::io::Error::last_os_error();
3847 if error.kind() == std::io::ErrorKind::Interrupted {
3848 continue;
3849 }
3850 return Err(error.into());
3851 }
3852 Ok(id)
3853}
3854
3855fn read_exact_from_fd(
3856 fd: i32,
3857 mut buf: &mut [u8],
3858 deadline: Instant,
3859 label: &str,
3860) -> AgentdResult<()> {
3861 while !buf.is_empty() {
3862 if !poll_fd_until(fd, libc::POLLIN, Some(deadline))? {
3863 return Err(AgentdError::ExecSession(format!(
3864 "timed out waiting for {label}"
3865 )));
3866 }
3867 match read_from_fd(fd, buf) {
3868 Ok(0) => {
3869 return Err(AgentdError::ExecSession(format!(
3870 "serial port closed while waiting for {label}"
3871 )));
3872 }
3873 Ok(read) => {
3874 let (_, remainder) = std::mem::take(&mut buf).split_at_mut(read);
3875 buf = remainder;
3876 }
3877 Err(error)
3878 if matches!(
3879 error.kind(),
3880 std::io::ErrorKind::Interrupted | std::io::ErrorKind::WouldBlock
3881 ) => {}
3882 Err(error) => return Err(error.into()),
3883 }
3884 }
3885 Ok(())
3886}
3887
3888fn init_ack_deadline() -> Instant {
3889 Instant::now() + std::time::Duration::from_secs(INIT_ACK_TIMEOUT_SECS)
3890}
3891
3892fn init_ack_timeout() -> AgentdError {
3893 AgentdError::ExecSession("timed out waiting for init ack".into())
3894}
3895
3896fn wait_for_init_ack(
3897 fd: i32,
3898 boot_console: &mut BootConsoleState,
3899 deadline: Instant,
3900) -> AgentdResult<()> {
3901 let msg = read_boot_message(fd, boot_console, Some(deadline), "init ack")?;
3902 if msg.t == MessageType::InitAck {
3903 let _: InitAck = msg
3904 .payload()
3905 .map_err(|e| AgentdError::ExecSession(format!("decode init ack payload: {e}")))?;
3906 return Ok(());
3907 }
3908
3909 Err(AgentdError::ExecSession(format!(
3910 "expected core.init.ack, got {}",
3911 msg.t.as_str()
3912 )))
3913}
3914
3915fn read_boot_message(
3916 fd: i32,
3917 state: &mut BootConsoleState,
3918 deadline: Option<Instant>,
3919 context: &str,
3920) -> AgentdResult<Message> {
3921 let mut read_buf = [0u8; 4096];
3922 loop {
3923 if let Some(msg) = codec::try_decode_from_buf(&mut state.input)
3924 .map_err(|e| AgentdError::ExecSession(format!("decode {context}: {e}")))?
3925 {
3926 return Ok(msg);
3927 }
3928 if state.input.len() > MAX_INPUT_BUF_SIZE {
3929 return Err(AgentdError::ExecSession(format!(
3930 "serial input buffer exceeded maximum size while waiting for {context}"
3931 )));
3932 }
3933 if !poll_fd_until(fd, libc::POLLIN, deadline)? {
3934 return Err(if context == "init ack" {
3935 init_ack_timeout()
3936 } else {
3937 AgentdError::ExecSession(format!("timed out waiting for {context}"))
3938 });
3939 }
3940 let n = match read_from_fd(fd, &mut read_buf) {
3941 Ok(n) => n,
3942 Err(error)
3943 if matches!(
3944 error.kind(),
3945 std::io::ErrorKind::Interrupted | std::io::ErrorKind::WouldBlock
3946 ) =>
3947 {
3948 continue;
3949 }
3950 Err(error) => return Err(error.into()),
3951 };
3952 if n == 0 {
3953 return Err(AgentdError::ExecSession(format!(
3954 "serial port closed while waiting for {context}"
3955 )));
3956 }
3957 state.input.extend_from_slice(&read_buf[..n]);
3958 }
3959}
3960
3961fn poll_fd_until(fd: i32, events: i16, deadline: Option<Instant>) -> AgentdResult<bool> {
3962 loop {
3963 let timeout_ms = match deadline {
3964 Some(deadline) => {
3965 let remaining = deadline.saturating_duration_since(Instant::now());
3966 if remaining.is_zero() {
3967 return Ok(false);
3968 }
3969 remaining.as_millis().clamp(1, i32::MAX as u128) as i32
3970 }
3971 None => -1,
3972 };
3973 let mut pfd = libc::pollfd {
3974 fd,
3975 events,
3976 revents: 0,
3977 };
3978 let ret = unsafe { libc::poll(&mut pfd, 1, timeout_ms) };
3979 if ret > 0 {
3980 return Ok(true);
3981 }
3982 if ret == 0 {
3983 return Ok(false);
3984 }
3985 let err = std::io::Error::last_os_error();
3986 if err.raw_os_error() == Some(libc::EINTR) {
3987 continue;
3988 }
3989 return Err(err.into());
3990 }
3991}
3992
3993fn read_from_fd(fd: i32, buf: &mut [u8]) -> std::io::Result<usize> {
3995 let n = unsafe { libc::read(fd, buf.as_mut_ptr() as *mut libc::c_void, buf.len()) };
3996 if n < 0 {
3997 Err(std::io::Error::last_os_error())
3998 } else {
3999 Ok(n as usize)
4000 }
4001}
4002
4003fn write_all_to_fd(fd: i32, mut buf: &[u8], deadline: Option<Instant>) -> AgentdResult<()> {
4004 while !buf.is_empty() {
4005 match write_to_fd(fd, buf) {
4006 Ok(0) => return Err(std::io::Error::from(std::io::ErrorKind::WriteZero).into()),
4007 Ok(n) => buf = &buf[n..],
4008 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
4009 Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
4010 if !poll_fd_until(fd, libc::POLLOUT, deadline)? {
4011 return Err(init_ack_timeout());
4012 }
4013 }
4014 Err(e) => return Err(e.into()),
4015 }
4016 }
4017
4018 Ok(())
4019}
4020
4021async fn flush_write_buf(fd: &AsyncFd<std::fs::File>, buf: &mut Vec<u8>) -> AgentdResult<()> {
4023 write_all_async_fd(fd, buf).await?;
4024 buf.clear();
4025 Ok(())
4026}
4027
4028async fn write_all_async_fd(fd: &AsyncFd<std::fs::File>, buf: &[u8]) -> AgentdResult<()> {
4030 let mut written = 0;
4031 while written < buf.len() {
4032 let mut guard = fd.writable().await?;
4033 match guard.try_io(|inner| write_to_fd(inner.get_ref().as_raw_fd(), &buf[written..])) {
4034 Ok(Ok(n)) => {
4035 if n == 0 {
4036 return Err(std::io::Error::from(std::io::ErrorKind::WriteZero).into());
4037 }
4038 written += n;
4039 }
4040 Ok(Err(e)) if e.kind() == std::io::ErrorKind::Interrupted => continue,
4041 Ok(Err(e)) => return Err(e.into()),
4042 Err(_would_block) => continue,
4043 }
4044 }
4045 Ok(())
4046}
4047
4048async fn write_bulk_record_async_fd(
4050 fd: &AsyncFd<std::fs::File>,
4051 record: &BulkRecord,
4052) -> AgentdResult<()> {
4053 let header = codec::encode_bulk_header(record)
4054 .map_err(|error| AgentdError::ExecSession(format!("encode bulk header: {error}")))?;
4055 write_bulk_parts_async_fd(fd, &header, &record.payload).await
4056}
4057
4058async fn write_incarnated_bulk_record_async_fd(
4060 fd: &AsyncFd<std::fs::File>,
4061 incarnation: ClientIncarnation,
4062 record: &BulkRecord,
4063) -> AgentdResult<()> {
4064 let public_header = codec::encode_bulk_header(record)
4065 .map_err(|error| AgentdError::ExecSession(format!("encode bulk header: {error}")))?;
4066 let mut header = [0u8; CLIENT_INCARNATION_SIZE + 4 + FRAME_HEADER_SIZE + BULK_HEADER_SIZE];
4067 header[..CLIENT_INCARNATION_SIZE].copy_from_slice(&incarnation);
4068 header[CLIENT_INCARNATION_SIZE..].copy_from_slice(&public_header);
4069 write_bulk_parts_async_fd(fd, &header, &record.payload).await
4070}
4071
4072async fn write_bulk_parts_async_fd(
4073 fd: &AsyncFd<std::fs::File>,
4074 header: &[u8],
4075 payload: &[u8],
4076) -> AgentdResult<()> {
4077 let mut header_offset = 0;
4078 let mut payload_offset = 0;
4079
4080 while header_offset < header.len() || payload_offset < payload.len() {
4081 let mut guard = fd.writable().await?;
4082 let result = guard.try_io(|inner| {
4083 write_vectored_to_fd(
4084 inner.get_ref().as_raw_fd(),
4085 &header[header_offset..],
4086 &payload[payload_offset..],
4087 )
4088 });
4089 let written = match result {
4090 Ok(Ok(0)) => {
4091 return Err(std::io::Error::from(std::io::ErrorKind::WriteZero).into());
4092 }
4093 Ok(Ok(written)) => written,
4094 Ok(Err(error)) if error.kind() == std::io::ErrorKind::Interrupted => continue,
4095 Ok(Err(error)) => return Err(error.into()),
4096 Err(_would_block) => continue,
4097 };
4098
4099 let header_remaining = header.len() - header_offset;
4100 if written < header_remaining {
4101 header_offset += written;
4102 } else {
4103 header_offset = header.len();
4104 payload_offset += written - header_remaining;
4105 }
4106 }
4107
4108 Ok(())
4109}
4110
4111fn write_to_fd(fd: i32, buf: &[u8]) -> std::io::Result<usize> {
4113 let n = unsafe { libc::write(fd, buf.as_ptr() as *const libc::c_void, buf.len()) };
4114 if n < 0 {
4115 Err(std::io::Error::last_os_error())
4116 } else {
4117 Ok(n as usize)
4118 }
4119}
4120
4121fn write_vectored_to_fd(fd: i32, header: &[u8], payload: &[u8]) -> std::io::Result<usize> {
4122 if header.is_empty() {
4123 return write_to_fd(fd, payload);
4124 }
4125
4126 let vectors = [
4127 libc::iovec {
4128 iov_base: header.as_ptr().cast_mut().cast(),
4129 iov_len: header.len(),
4130 },
4131 libc::iovec {
4132 iov_base: payload.as_ptr().cast_mut().cast(),
4133 iov_len: payload.len(),
4134 },
4135 ];
4136 let vector_count = if payload.is_empty() { 1 } else { 2 };
4137 let written = unsafe { libc::writev(fd, vectors.as_ptr(), vector_count) };
4138 if written < 0 {
4139 Err(std::io::Error::last_os_error())
4140 } else {
4141 Ok(written as usize)
4142 }
4143}
4144
4145fn request_guest_poweroff() -> AgentdResult<()> {
4146 if crate::handoff::is_pid_1() {
4147 crate::teardown::teardown_filesystems(true);
4150 let ret = unsafe { libc::reboot(libc::RB_POWER_OFF) };
4151 if ret != 0 {
4152 return Err(std::io::Error::last_os_error().into());
4153 }
4154 return Ok(());
4155 }
4156
4157 unsafe {
4158 libc::sync();
4159 }
4160
4161 crate::handoff::signal_init_shutdown()
4164}
4165
4166#[cfg(test)]
4171mod tests {
4172 use super::*;
4173 use bytes::Bytes;
4174 use microsandbox_protocol::message::PROTOCOL_VERSION;
4175
4176 #[allow(clippy::too_many_arguments)]
4177 async fn handle_message(
4178 msg: Message,
4179 state: &mut AgentState,
4180 activity: &mut ActivityTracker,
4181 sender: &mut SessionOutputSender,
4182 out_buf: &mut Vec<u8>,
4183 config: &AgentdConfig,
4184 workload: &mut WorkloadLatch,
4185 heartbeat: &heartbeat::HeartbeatControl,
4186 ) -> AgentdResult<()> {
4187 handle_message_with_charge(
4188 msg, state, activity, sender, out_buf, config, workload, heartbeat, None,
4189 )
4190 .await
4191 }
4192
4193 fn decode_reply_skipping_credit(bytes: &mut BytesMut) -> Message {
4194 loop {
4195 let Some(DecodedFrame::Control(reply)) =
4196 codec::try_decode_frame_from_bytes(bytes).unwrap()
4197 else {
4198 panic!("missing control reply");
4199 };
4200 if reply.t != MessageType::WorkloadTransportCredit {
4201 return reply;
4202 }
4203 assert_eq!(reply.id, u32::MAX);
4204 }
4205 }
4206
4207 #[test]
4208 fn admitted_transport_window_fits_each_filesystem_input_queue() {
4209 assert!(
4212 WORKLOAD_TRANSPORT_CONTROL_FRAMES + WORKLOAD_TRANSPORT_BULK_FRAMES
4213 <= FS_BULK_INPUT_ITEM_CAPACITY as u64
4214 );
4215 assert!(WORKLOAD_TRANSPORT_CONTROL_BYTES <= BULK_INPUT_BYTE_CAPACITY as u64);
4216 assert!(WORKLOAD_TRANSPORT_BULK_BYTES <= BULK_INPUT_BYTE_CAPACITY as u64);
4217 }
4218
4219 #[test]
4220 fn private_admission_and_credit_batching_use_the_complete_wire_contract() {
4221 let mut request = Message::with_payload(
4222 MessageType::WorkloadFreeze,
4223 u32::MAX,
4224 &WorkloadFreeze {
4225 external_mount_tags: Vec::new(),
4226 attempt_id: "cut".into(),
4227 host_input: Default::default(),
4228 },
4229 )
4230 .unwrap();
4231 assert!(private_lifecycle_message(&request));
4232 request.id = 0;
4233 assert!(!private_lifecycle_message(&request));
4234 request.id = u32::MAX;
4235 request.t = MessageType::ClockSync;
4236 assert!(!private_lifecycle_message(&request));
4237 let record = BulkRecord {
4238 id: 1,
4239 kind: BulkKind::Tcp,
4240 flow: BulkFlow::HostToGuest,
4241 offset: 0,
4242 payload: Bytes::from_static(b"payload"),
4243 };
4244 let mut encoded = codec::encode_bulk_header(&record).unwrap().to_vec();
4245 encoded.extend_from_slice(&record.payload);
4246 assert_eq!(bulk_wire_bytes(&record, false), encoded.len());
4247 assert_eq!(
4248 bulk_wire_bytes(&record, true),
4249 encoded.len() + CLIENT_INCARNATION_SIZE
4250 );
4251 let previous = initial_input_credit();
4252 let mut current = previous;
4253 current.control_frames += 15;
4254 assert!(!input_credit_update_due(&previous, ¤t));
4255 current.control_frames += 1;
4256 assert!(input_credit_update_due(&previous, ¤t));
4257 let mut credit = Vec::new();
4258 encode_input_credit(current, &mut credit).unwrap();
4259 let mut bytes = BytesMut::from(credit.as_slice());
4260 let Some(DecodedFrame::Control(reply)) =
4261 codec::try_decode_frame_from_bytes(&mut bytes).unwrap()
4262 else {
4263 panic!()
4264 };
4265 assert_eq!(reply.id, u32::MAX);
4266 assert_eq!(reply.payload::<WorkloadTransportCredit>().unwrap(), current);
4267 assert!(!guest_message_refreshes_idle_timer(&reply.t));
4268 }
4269
4270 #[tokio::test]
4271 async fn transport_cut_waits_for_decoded_prefix_and_abort_is_owned_and_repeatable() {
4272 use microsandbox_protocol::core::WorkloadThawMode::{Continue, Restore};
4273 let mut state = AgentState::default();
4274 let (mut sender, _output) = SessionOutputSender::channel();
4275 let mut activity = ActivityTracker::new();
4276 let config = AgentdConfig {
4277 user: None,
4278 security_profile: Default::default(),
4279 default_cwd: None,
4280 default_env: Vec::new(),
4281 };
4282 let heartbeat = heartbeat::HeartbeatControl::default();
4283 let mut workload = crate::workload::tests::fake_latch();
4284 let target = WorkloadTransportPosition {
4285 bulk_bytes: 80,
4286 bulk_frames: 1,
4287 ..Default::default()
4288 };
4289 let freeze = || {
4290 Message::with_payload(
4291 MessageType::WorkloadFreeze,
4292 u32::MAX,
4293 &WorkloadFreeze {
4294 external_mount_tags: Vec::new(),
4295 attempt_id: "prefix".into(),
4296 host_input: target,
4297 },
4298 )
4299 .unwrap()
4300 };
4301 for _ in 0..2 {
4302 let mut out = Vec::new();
4303 handle_message(
4304 freeze(),
4305 &mut state,
4306 &mut activity,
4307 &mut sender,
4308 &mut out,
4309 &config,
4310 &mut workload,
4311 &heartbeat,
4312 )
4313 .await
4314 .unwrap();
4315 assert!(out.is_empty());
4316 assert!(!workload.is_frozen());
4317 assert!(state.pending_freeze.is_some());
4318 }
4319 for (attempt, mode, expected) in [
4320 ("other", Continue, MessageType::CoreError),
4321 ("prefix", Restore, MessageType::CoreError),
4322 ("prefix", Continue, MessageType::WorkloadThawed),
4323 ("prefix", Continue, MessageType::WorkloadThawed),
4324 ] {
4325 let thaw = Message::with_payload(
4326 MessageType::WorkloadThaw,
4327 u32::MAX,
4328 &WorkloadThaw {
4329 attempt_id: attempt.into(),
4330 mode,
4331 },
4332 )
4333 .unwrap();
4334 let aborts_pending =
4335 expected == MessageType::WorkloadThawed && state.pending_freeze.is_some();
4336 let mut out = Vec::new();
4337 handle_message(
4338 thaw,
4339 &mut state,
4340 &mut activity,
4341 &mut sender,
4342 &mut out,
4343 &config,
4344 &mut workload,
4345 &heartbeat,
4346 )
4347 .await
4348 .unwrap();
4349 let mut bytes = BytesMut::from(out.as_slice());
4350 if aborts_pending {
4351 let cancelled = decode_reply_skipping_credit(&mut bytes);
4352 let error = cancelled.payload::<CoreError>().unwrap();
4353 assert_eq!(cancelled.id, u32::MAX);
4354 assert_eq!(
4355 error.offending_type.as_deref(),
4356 Some(MessageType::WorkloadFreeze.as_str())
4357 );
4358 assert_eq!(error.workload_failure.unwrap().attempt_id, "prefix");
4359 }
4360 assert_eq!(decode_reply_skipping_credit(&mut bytes).t, expected);
4361 assert_eq!(
4362 state.pending_freeze.is_some(),
4363 expected == MessageType::CoreError
4364 );
4365 }
4366 let charge = state.input_window.admit(InputLane::Bulk, 80).unwrap();
4367 for _ in 0..2 {
4368 let mut out = Vec::new();
4369 handle_message(
4370 freeze(),
4371 &mut state,
4372 &mut activity,
4373 &mut sender,
4374 &mut out,
4375 &config,
4376 &mut workload,
4377 &heartbeat,
4378 )
4379 .await
4380 .unwrap();
4381 let reply = decode_reply_skipping_credit(&mut BytesMut::from(out.as_slice()));
4382 let frozen = reply.payload::<WorkloadFrozen>().unwrap();
4383 assert_eq!(reply.t, MessageType::WorkloadFrozen);
4384 assert!(workload.is_frozen());
4385 assert!(state.output_parked);
4386 assert_eq!(
4387 frozen.input_credit,
4388 initial_input_credit(),
4389 "decode is not consumption"
4390 );
4391 assert_eq!(state.input_window.position(), target);
4392 }
4393 drop(charge);
4394 }
4395
4396 #[tokio::test]
4397 async fn mixed_primary_data_cut_waits_for_the_complete_dedicated_prefix() {
4398 let mut state = AgentState::default();
4399 let (mut sender, _output) = SessionOutputSender::channel();
4400 let mut activity = ActivityTracker::new();
4401 let config = AgentdConfig {
4402 user: None,
4403 security_profile: Default::default(),
4404 default_cwd: None,
4405 default_env: Vec::new(),
4406 };
4407 let heartbeat = heartbeat::HeartbeatControl::default();
4408 let mut workload = crate::workload::tests::fake_latch();
4409
4410 let mut charges = Vec::new();
4413 for data in [vec![0x31; 1024], Vec::new()] {
4414 let message =
4415 Message::with_payload(MessageType::ExecStdin, 1, &ExecStdin { data }).unwrap();
4416 assert!(message.t.uses_workload_data_credit());
4417 let mut wire = Vec::new();
4418 codec::encode_to_buf(&message, &mut wire).unwrap();
4419 charges.push(
4420 state
4421 .input_window
4422 .admit(InputLane::Bulk, wire.len())
4423 .unwrap(),
4424 );
4425 }
4426 let primary_position = state.input_window.position();
4427 assert_eq!(primary_position.control_bytes, 0);
4428 assert_eq!(primary_position.control_frames, 0);
4429 assert_eq!(primary_position.bulk_frames, 2);
4430
4431 let incarnation = [0x43; CLIENT_INCARNATION_SIZE];
4432 let record = BulkRecord {
4433 id: 2,
4434 kind: BulkKind::Filesystem,
4435 flow: BulkFlow::HostToGuest,
4436 offset: 0,
4437 payload: Bytes::from(vec![0x52; 512]),
4438 };
4439 let mut dedicated_wire = incarnation.to_vec();
4440 codec::encode_bulk_to_buf(&record, &mut dedicated_wire).unwrap();
4441 let target = WorkloadTransportPosition {
4442 bulk_bytes: primary_position.bulk_bytes + dedicated_wire.len() as u64,
4443 bulk_frames: primary_position.bulk_frames + 1,
4444 ..primary_position
4445 };
4446 let freeze = Message::with_payload(
4447 MessageType::WorkloadFreeze,
4448 u32::MAX,
4449 &WorkloadFreeze {
4450 external_mount_tags: Vec::new(),
4451 attempt_id: "mixed-prefix".into(),
4452 host_input: target,
4453 },
4454 )
4455 .unwrap();
4456
4457 let mut dedicated_input = BytesMut::new();
4458 for prefix in [
4459 &dedicated_wire[..0],
4460 &dedicated_wire[..dedicated_wire.len() - 1],
4461 ] {
4462 dedicated_input.extend_from_slice(prefix);
4463 assert!(
4464 try_decode_incarnated_bulk_from_bytes(&mut dedicated_input)
4465 .unwrap()
4466 .is_none()
4467 );
4468 let mut out = Vec::new();
4469 handle_message(
4470 freeze.clone(),
4471 &mut state,
4472 &mut activity,
4473 &mut sender,
4474 &mut out,
4475 &config,
4476 &mut workload,
4477 &heartbeat,
4478 )
4479 .await
4480 .unwrap();
4481 assert!(
4482 out.is_empty(),
4483 "primary data cannot cover missing dedicated bytes"
4484 );
4485 assert!(!workload.is_frozen());
4486 assert!(!state.output_parked);
4487 assert!(state.pending_freeze.is_some());
4488 assert_eq!(state.input_window.position(), primary_position);
4489 }
4490
4491 dedicated_input.extend_from_slice(&dedicated_wire[dedicated_wire.len() - 1..]);
4492 let decoded = try_decode_incarnated_bulk_from_bytes(&mut dedicated_input)
4493 .unwrap()
4494 .unwrap();
4495 assert_eq!(decoded.incarnation, incarnation);
4496 assert_eq!(decoded.record, record);
4497 assert!(dedicated_input.is_empty());
4498 charges.push(
4499 state
4500 .input_window
4501 .admit(InputLane::Bulk, bulk_wire_bytes(&decoded.record, true))
4502 .unwrap(),
4503 );
4504 assert_eq!(state.input_window.position(), target);
4505
4506 let pending = state.pending_freeze.take().unwrap();
4509 let mut out = Vec::new();
4510 handle_message(
4511 pending,
4512 &mut state,
4513 &mut activity,
4514 &mut sender,
4515 &mut out,
4516 &config,
4517 &mut workload,
4518 &heartbeat,
4519 )
4520 .await
4521 .unwrap();
4522 let mut bytes = BytesMut::from(out.as_slice());
4523 let reply = decode_reply_skipping_credit(&mut bytes);
4524 let frozen = reply.payload::<WorkloadFrozen>().unwrap();
4525 assert_eq!(reply.t, MessageType::WorkloadFrozen);
4526 assert_eq!(reply.id, u32::MAX);
4527 assert_eq!(frozen.attempt_id, "mixed-prefix");
4528 assert_eq!(frozen.input_credit, initial_input_credit());
4529 assert_eq!(state.frozen_host_input, Some(target));
4530 assert!(workload.is_frozen());
4531 assert!(state.output_parked);
4532 assert!(state.pending_freeze.is_none());
4533 assert!(bytes.is_empty());
4534 drop(charges);
4535 }
4536
4537 #[tokio::test]
4538 async fn late_tcp_credit_preserves_raw_tail_before_and_after_terminal_retirement() {
4539 use std::collections::hash_map::DefaultHasher;
4540 use std::hash::Hasher;
4541
4542 use microsandbox_protocol::bulk::BulkOffer;
4543 use tokio::io::{AsyncReadExt, AsyncWriteExt};
4544 use tokio::net::TcpListener;
4545
4546 tokio::time::timeout(Duration::from_secs(10), async {
4547 let mut state = AgentState::default();
4548 let incarnation = [0x57; CLIENT_INCARNATION_SIZE];
4549 establish_relay_client(
4550 &mut state,
4551 RelayClientConnected {
4552 id_start: 1,
4553 id_end_exclusive: microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP,
4554 incarnation,
4555 },
4556 )
4557 .unwrap();
4558 let (mut sender, mut control, mut bulk, mut scheduler_commands) =
4561 SessionOutputSender::split_channel();
4562 let producer = sender.with_incarnation(Some(incarnation));
4563 let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
4564 state.tcp_sessions.insert(
4565 1,
4566 TcpSession::open(
4567 1,
4568 TcpConnect {
4569 host: "127.0.0.1".into(),
4570 port: listener.local_addr().unwrap().port(),
4571 bulk: Some(BulkOffer::tcp()),
4572 },
4573 &producer,
4574 ),
4575 );
4576 let (mut peer, _) = listener.accept().await.unwrap();
4577 for expected in [MessageType::TcpConnected, MessageType::BulkAccepted] {
4578 let envelope = control.recv().await.unwrap();
4579 let SessionOutput::Raw(mut output) = envelope.output else {
4580 panic!()
4581 };
4582 assert_eq!(
4583 codec::try_decode_from_buf(&mut output.frame)
4584 .unwrap()
4585 .unwrap()
4586 .t,
4587 expected
4588 );
4589 }
4590 state
4591 .tcp_sessions
4592 .get(&1)
4593 .unwrap()
4594 .finish_bulk(BulkFinish {
4595 kind: BulkKind::Tcp,
4596 flow: BulkFlow::HostToGuest,
4597 final_offset: 0,
4598 })
4599 .await
4600 .unwrap();
4601 assert_eq!(peer.read(&mut [0]).await.unwrap(), 0);
4602 let payload = Bytes::from(
4603 (0..6 * 1024 * 1024)
4604 .map(|index| (index % 251) as u8)
4605 .collect::<Vec<_>>(),
4606 );
4607 let peer_payload = payload.clone();
4608 let peer_task = tokio::spawn(async move {
4609 peer.write_all(&peer_payload).await.unwrap();
4610 peer.shutdown().await.unwrap();
4611 });
4612 let mut received = Vec::new();
4613 let mut consumed = 0;
4614 while consumed < 4 * 1024 * 1024 {
4615 let envelope = bulk.recv().await.unwrap();
4616 let SessionOutput::Bulk(output) = &envelope.output else {
4617 panic!()
4618 };
4619 consumed += output.record.payload.len();
4620 received.push(envelope);
4621 }
4622 peer_task.await.unwrap();
4623 while !state.tcp_sessions.get(&1).unwrap().is_finished() {
4624 tokio::task::yield_now().await;
4625 }
4626 assert!(consumed < payload.len());
4627 assert!(
4628 !bulk.is_empty(),
4629 "the dedicated output tail must still be queued"
4630 );
4631 let credit = BulkCredit {
4632 kind: BulkKind::Tcp,
4633 flow: BulkFlow::GuestToHost,
4634 consumed_offset: consumed as u64,
4635 credit_limit: consumed as u64 + DEFAULT_BULK_WINDOW,
4636 };
4637 let mut activity = ActivityTracker::new();
4638 let config = AgentdConfig {
4639 user: None,
4640 security_profile: Default::default(),
4641 default_cwd: None,
4642 default_env: Vec::new(),
4643 };
4644 let mut workload = crate::workload::tests::fake_latch();
4645 let heartbeat = heartbeat::HeartbeatControl::default();
4646 for retired in [false, true] {
4647 if retired {
4648 for expected in [MessageType::BulkFinish, MessageType::TcpClosed] {
4649 let envelope = control.try_recv().unwrap();
4650 let SessionOutput::Raw(mut output) = envelope.output else {
4651 panic!()
4652 };
4653 let message = codec::try_decode_from_buf(&mut output.frame)
4654 .unwrap()
4655 .unwrap();
4656 assert_eq!(message.t, expected);
4657 if expected == MessageType::BulkFinish {
4658 assert_eq!(
4659 message.payload::<BulkFinish>().unwrap().final_offset,
4660 payload.len() as u64
4661 );
4662 } else {
4663 assert_ne!(
4664 message.flags & microsandbox_protocol::message::FLAG_TERMINAL,
4665 0
4666 );
4667 assert!(matches!(output.completion, Some(RawSessionCompletion::Tcp)));
4668 complete_raw_session(
4669 1,
4670 output.completion,
4671 &mut state.read_sessions,
4672 &mut state.tcp_sessions,
4673 );
4674 clear_bulk_receive_state(&mut state, 1);
4675 }
4676 }
4677 }
4678 assert_eq!(state.tcp_sessions.contains_key(&1), !retired);
4679 let mut out = Vec::new();
4680 handle_message(
4681 Message::with_payload(MessageType::BulkCredit, 1, &credit).unwrap(),
4682 &mut state,
4683 &mut activity,
4684 &mut sender,
4685 &mut out,
4686 &config,
4687 &mut workload,
4688 &heartbeat,
4689 )
4690 .await
4691 .unwrap();
4692 assert!(
4693 out.is_empty(),
4694 "late credit must not emit cancellation or another terminal"
4695 );
4696 assert!(
4697 matches!(
4698 scheduler_commands.try_recv(),
4699 Err(tokio::sync::mpsc::error::TryRecvError::Empty)
4700 ),
4701 "late credit must not purge raw output"
4702 );
4703 assert!(!bulk.is_empty());
4704 }
4705 assert!(matches!(
4706 control.try_recv(),
4707 Err(tokio::sync::mpsc::error::TryRecvError::Empty)
4708 ));
4709 while let Ok(envelope) = bulk.try_recv() {
4710 received.push(envelope);
4711 }
4712 let mut offset = 0;
4713 let mut actual_hash = DefaultHasher::new();
4714 let mut expected_hash = DefaultHasher::new();
4715 expected_hash.write(&payload);
4716 for envelope in received {
4717 assert_eq!(envelope.id, 1);
4718 assert_eq!(envelope.incarnation, Some(incarnation));
4719 let SessionOutput::Bulk(output) = envelope.output else {
4720 panic!()
4721 };
4722 let record = output.record;
4723 assert_eq!(record.offset, offset as u64);
4724 let end = offset + record.payload.len();
4725 assert_eq!(record.payload.as_ref(), &payload[offset..end]);
4726 actual_hash.write(&record.payload);
4727 offset = end;
4728 }
4729 assert_eq!(offset, payload.len());
4730 assert_eq!(actual_hash.finish(), expected_hash.finish());
4731 })
4732 .await
4733 .expect("late-credit TCP tail did not complete");
4734 }
4735
4736 #[tokio::test]
4737 async fn dedicated_bulk_park_finishes_one_record_and_retains_the_next() {
4738 use std::os::fd::OwnedFd;
4739 use tokio::io::AsyncReadExt;
4740
4741 let (writer, reader) = std::os::unix::net::UnixStream::pair().unwrap();
4742 writer.set_nonblocking(true).unwrap();
4743 reader.set_nonblocking(true).unwrap();
4744 let send_buffer: libc::c_int = 4096;
4745 assert_eq!(
4746 unsafe {
4747 libc::setsockopt(
4748 writer.as_raw_fd(),
4749 libc::SOL_SOCKET,
4750 libc::SO_SNDBUF,
4751 (&send_buffer as *const libc::c_int).cast(),
4752 std::mem::size_of_val(&send_buffer) as libc::socklen_t,
4753 )
4754 },
4755 0
4756 );
4757 let file = File::from(OwnedFd::from(writer));
4758 let mut reader = tokio::net::UnixStream::from_std(reader).unwrap();
4759 let (sender, _control, bulk, commands) = SessionOutputSender::split_channel();
4760 let (activity, _activities) = tokio::sync::mpsc::channel(8);
4761 let writer = tokio::spawn(bulk_writer_task(file, bulk, commands, activity));
4762 let sender = sender.with_incarnation(Some([0x62; CLIENT_INCARNATION_SIZE]));
4763 let first = BulkRecord {
4764 id: 1,
4765 kind: BulkKind::Filesystem,
4766 flow: BulkFlow::GuestToHost,
4767 offset: 0,
4768 payload: Bytes::from(vec![0x6a; 1024 * 1024]),
4769 };
4770 let first_len = bulk_wire_bytes(&first, true);
4771 let second = BulkRecord {
4772 offset: first.payload.len() as u64,
4773 payload: Bytes::from_static(b"next"),
4774 ..first.clone()
4775 };
4776 let second_len = bulk_wire_bytes(&second, true);
4777 for record in [first, second] {
4778 assert!(
4779 sender
4780 .send(
4781 1,
4782 SessionOutput::Bulk(crate::session::BulkSessionOutput::new(
4783 record,
4784 RawActivity::default()
4785 ))
4786 )
4787 .await
4788 );
4789 }
4790 let mut bytes = vec![0; first_len];
4791 time::timeout(Duration::from_secs(5), reader.read_exact(&mut bytes[..37]))
4792 .await
4793 .unwrap()
4794 .unwrap();
4795 let park_sender = sender.clone();
4796 let mut park = tokio::spawn(async move { park_sender.park_bulk_output().await });
4797 assert!(
4798 time::timeout(Duration::from_millis(20), &mut park)
4799 .await
4800 .is_err(),
4801 "park cut a partial record"
4802 );
4803 time::timeout(Duration::from_secs(5), reader.read_exact(&mut bytes[37..]))
4804 .await
4805 .unwrap()
4806 .unwrap();
4807 assert_eq!(
4808 time::timeout(Duration::from_secs(5), park)
4809 .await
4810 .unwrap()
4811 .unwrap()
4812 .unwrap(),
4813 first_len as u64
4814 );
4815 let mut decoded = BytesMut::from(bytes.as_slice());
4816 let frame = try_decode_incarnated_bulk_from_bytes(&mut decoded)
4817 .unwrap()
4818 .unwrap();
4819 assert_eq!(frame.record.payload, Bytes::from(vec![0x6a; 1024 * 1024]));
4820 assert!(decoded.is_empty());
4821 assert!(
4822 time::timeout(Duration::from_millis(20), reader.read(&mut [0; 1]))
4823 .await
4824 .is_err()
4825 );
4826 sender.resume_bulk_output().await.unwrap();
4827 let mut bytes = vec![0; second_len];
4828 time::timeout(Duration::from_secs(5), reader.read_exact(&mut bytes))
4829 .await
4830 .unwrap()
4831 .unwrap();
4832 assert_eq!(
4833 sender.park_bulk_output().await.unwrap(),
4834 (first_len + second_len) as u64
4835 );
4836 writer.abort();
4837 let _ = writer.await;
4838 }
4839
4840 #[test]
4841 fn init_failure_waits_for_failure_ack_after_late_success_ack() {
4842 use std::io::{Read, Write};
4843 use std::os::fd::OwnedFd;
4844 use std::os::unix::net::UnixStream;
4845
4846 for ack_failure in [false, true] {
4847 let (guest, mut host) = UnixStream::pair().unwrap();
4848 host.set_read_timeout(Some(std::time::Duration::from_secs(5)))
4849 .unwrap();
4850 let guest = File::from(OwnedFd::from(guest));
4851 let worker = std::thread::spawn(move || {
4852 report_init_failure(
4853 &guest,
4854 &mut BootConsoleState::default(),
4855 "missing guest user",
4856 &AgentdError::UserNotFound("iggy".into()),
4857 )
4858 });
4859 let mut header = [0u8; 4];
4860 host.read_exact(&mut header).unwrap();
4861 let mut frame = vec![0; u32::from_be_bytes(header) as usize];
4862 host.read_exact(&mut frame).unwrap();
4863 let mut input = header.to_vec();
4864 input.extend_from_slice(&frame);
4865 let report = codec::try_decode_from_buf(&mut input).unwrap().unwrap();
4866 assert_eq!(report.t, MessageType::CoreError);
4867 assert_eq!(
4868 report.payload::<CoreError>().unwrap().init_failure,
4869 Some(InitFailureReason::UserNotFound)
4870 );
4871 assert_eq!(
4872 report.payload::<CoreError>().unwrap().message,
4873 "missing guest user"
4874 );
4875
4876 let mut acks = Vec::new();
4877 for failure in [false].into_iter().chain(ack_failure.then_some(true)) {
4878 codec::encode_to_buf(
4879 &Message::with_payload(MessageType::InitAck, 0, &InitAck { failure }).unwrap(),
4880 &mut acks,
4881 )
4882 .unwrap();
4883 }
4884 host.write_all(&acks).unwrap();
4885 host.shutdown(std::net::Shutdown::Write).unwrap();
4886 assert_eq!(worker.join().unwrap().is_ok(), ack_failure);
4887 }
4888 }
4889
4890 #[test]
4891 fn init_failure_waits_for_slow_supporting_host_but_bounds_legacy_wait() {
4892 use std::io::Write;
4893 use std::os::fd::OwnedFd;
4894 use std::os::unix::net::UnixStream;
4895 use std::sync::mpsc;
4896
4897 for supported in [false, true] {
4898 let (guest, mut host) = UnixStream::pair().unwrap();
4899 let guest = File::from(OwnedFd::from(guest));
4900 let mut bootstrap = Vec::new();
4901 codec::encode_to_buf(
4902 &Message::with_payload(
4903 MessageType::Bootstrap,
4904 0,
4905 &GuestBootstrap {
4906 init_failure_ack: supported,
4907 ..Default::default()
4908 },
4909 )
4910 .unwrap(),
4911 &mut bootstrap,
4912 )
4913 .unwrap();
4914 host.write_all(&bootstrap).unwrap();
4915 let (done_tx, done_rx) = mpsc::channel();
4916 let worker = std::thread::spawn(move || {
4917 let (_, mut state) = receive_bootstrap(&guest).unwrap();
4918 let result = report_init_failure(
4919 &guest,
4920 &mut state,
4921 "missing user",
4922 &AgentdError::UserNotFound("iggy".into()),
4923 );
4924 done_tx.send(result).unwrap();
4925 });
4926 let report = read_boot_message(
4927 host.as_raw_fd(),
4928 &mut BootConsoleState::default(),
4929 Some(Instant::now() + std::time::Duration::from_secs(5)),
4930 "startup error",
4931 )
4932 .unwrap();
4933 assert_eq!(report.t, MessageType::CoreError);
4934
4935 let result = done_rx.recv_timeout(std::time::Duration::from_secs(3));
4936 if supported {
4937 assert!(
4938 matches!(result, Err(mpsc::RecvTimeoutError::Timeout)),
4939 "supporting host must be allowed to save beyond the legacy timeout"
4940 );
4941 let mut ack = Vec::new();
4942 codec::encode_to_buf(
4943 &Message::with_payload(MessageType::InitAck, 0, &InitAck { failure: true })
4944 .unwrap(),
4945 &mut ack,
4946 )
4947 .unwrap();
4948 host.write_all(&ack).unwrap();
4949 done_rx
4950 .recv_timeout(std::time::Duration::from_secs(5))
4951 .unwrap()
4952 .unwrap();
4953 } else {
4954 assert!(result.expect("legacy host wait must be bounded").is_err());
4955 }
4956 worker.join().unwrap();
4957 }
4958 }
4959
4960 #[test]
4961 fn coalesced_bootstrap_and_init_ack_retain_the_second_frame() {
4962 let bootstrap = GuestBootstrap::default();
4963 let bootstrap_message =
4964 Message::with_payload(MessageType::Bootstrap, 0, &bootstrap).unwrap();
4965 let ack_message =
4966 Message::with_payload(MessageType::InitAck, 0, &InitAck::default()).unwrap();
4967 let mut state = BootConsoleState::default();
4968 codec::encode_to_buf(&bootstrap_message, &mut state.input).unwrap();
4969 codec::encode_to_buf(&ack_message, &mut state.input).unwrap();
4970
4971 let decoded = read_boot_message(
4972 -1,
4973 &mut state,
4974 Some(Instant::now() + std::time::Duration::from_secs(1)),
4975 "guest bootstrap",
4976 )
4977 .unwrap();
4978 assert_eq!(decode_bootstrap_message(decoded).unwrap(), bootstrap);
4979 assert!(
4980 !state.input.is_empty(),
4981 "init ack frame should remain buffered"
4982 );
4983
4984 wait_for_init_ack(
4985 -1,
4986 &mut state,
4987 Instant::now() + std::time::Duration::from_secs(1),
4988 )
4989 .unwrap();
4990 assert!(state.input.is_empty());
4991 }
4992
4993 #[test]
4994 fn bootstrap_rejects_non_control_correlation_fields() {
4995 let mut message =
4996 Message::with_payload(MessageType::Bootstrap, 1, &GuestBootstrap::default()).unwrap();
4997 message.flags = 1;
4998
4999 let error = decode_bootstrap_message(message).unwrap_err();
5000 assert!(error.to_string().contains("requires id=0 and flags=0"));
5001 }
5002
5003 #[test]
5004 fn bootstrap_rejects_wrong_first_message_type() {
5005 let message = Message::with_payload(MessageType::Ping, 0, &Ping {}).unwrap();
5006
5007 let error = decode_bootstrap_message(message).unwrap_err();
5008 assert!(error.to_string().contains("expected core.bootstrap"));
5009 }
5010
5011 #[test]
5012 fn bootstrap_rejects_older_protocol_generation() {
5013 let mut message =
5014 Message::with_payload(MessageType::Bootstrap, 0, &GuestBootstrap::default()).unwrap();
5015 let min_version = MessageType::Bootstrap.min_protocol_version();
5016 assert!(min_version > 0);
5017 message.v = min_version - 1;
5018
5019 let error = decode_bootstrap_message(message).unwrap_err();
5020 assert!(error.to_string().contains("or newer"));
5021 }
5022
5023 #[test]
5024 fn bootstrap_accepts_minimum_supported_protocol_generation() {
5025 let mut message =
5026 Message::with_payload(MessageType::Bootstrap, 0, &GuestBootstrap::default()).unwrap();
5027 message.v = MessageType::Bootstrap.min_protocol_version();
5028
5029 assert_eq!(
5030 decode_bootstrap_message(message).unwrap(),
5031 GuestBootstrap::default()
5032 );
5033 }
5034
5035 #[test]
5036 fn bootstrap_accepts_newer_additive_protocol_generation() {
5037 let mut message =
5038 Message::with_payload(MessageType::Bootstrap, 0, &GuestBootstrap::default()).unwrap();
5039 message.v = PROTOCOL_VERSION + 1;
5040
5041 assert_eq!(
5042 decode_bootstrap_message(message).unwrap(),
5043 GuestBootstrap::default()
5044 );
5045 }
5046
5047 #[test]
5048 fn bootstrap_rejects_malformed_payload() {
5049 let message = Message::new(MessageType::Bootstrap, 0, vec![0xff]);
5050
5051 let error = decode_bootstrap_message(message).unwrap_err();
5052 assert!(error.to_string().contains("decode guest bootstrap payload"));
5053 }
5054 #[test]
5055 fn dual_port_cmdline_hint_requires_an_exact_argument() {
5056 assert!(cmdline_requests_dual_port(
5057 "console=hvc0 microsandbox.agent_transport=dual-port-v1 quiet"
5058 ));
5059 assert!(!cmdline_requests_dual_port(
5060 "microsandbox.agent_transport=dual-port-v10"
5061 ));
5062 assert!(!cmdline_requests_dual_port(
5063 "prefix=microsandbox.agent_transport=dual-port-v1"
5064 ));
5065 }
5066
5067 #[test]
5068 fn primary_read_buffer_tracks_whether_it_carries_bulk() {
5069 assert_eq!(
5070 primary_serial_read_buf_size(false),
5071 COMBINED_SERIAL_READ_BUF_SIZE
5072 );
5073 assert_eq!(
5074 primary_serial_read_buf_size(true),
5075 CONTROL_SERIAL_READ_BUF_SIZE
5076 );
5077 }
5078
5079 #[test]
5080 fn bulk_reader_turn_is_bounded_by_records_or_payload_bytes() {
5081 assert!(!bulk_reader_turn_exhausted(1, 1));
5082 assert!(bulk_reader_turn_exhausted(
5083 BULK_READER_MAX_RECORDS_PER_TURN,
5084 1
5085 ));
5086 assert!(bulk_reader_turn_exhausted(
5087 1,
5088 BULK_READER_MAX_BYTES_PER_TURN
5089 ));
5090 }
5091
5092 #[test]
5093 fn agent_read_budget_counts_bytes_and_calls_independently() {
5094 let mut bytes = AgentReadBudget::default();
5095 bytes.record_read(AGENT_READ_QUANTUM_BYTES - 1);
5096 assert!(!bytes.exhausted());
5097 bytes.record_read(1);
5098 assert!(bytes.exhausted());
5099
5100 let mut calls = AgentReadBudget::default();
5101 for _ in 1..AGENT_READ_QUANTUM_CALLS {
5102 assert!(calls.read_fd(-1, &mut [0]).is_err());
5103 assert!(!calls.exhausted());
5104 }
5105 assert!(calls.read_fd(-1, &mut [0]).is_err());
5106 assert!(calls.exhausted(), "failed reads also bound a retry loop");
5107 assert_eq!(calls.bytes, 0);
5108
5109 let mut large = AgentReadBudget::default();
5110 large.record_read(BULK_SERIAL_READ_BUF_SIZE);
5111 assert!(
5112 large.exhausted(),
5113 "a large read is not a smaller wire record"
5114 );
5115 assert_eq!(large.bytes, BULK_SERIAL_READ_BUF_SIZE);
5116 }
5117
5118 #[tokio::test(flavor = "current_thread")]
5119 async fn agent_read_budget_services_driver_during_partial_bulk_records() {
5120 assert!(!observe_driver_during_partial_bulk(false).await);
5123 assert!(observe_driver_during_partial_bulk(true).await);
5124 }
5125
5126 async fn observe_driver_during_partial_bulk(yield_at_boundary: bool) -> bool {
5127 use std::io::Write;
5128 use std::os::fd::OwnedFd;
5129 use std::os::unix::net::UnixStream;
5130 use std::sync::atomic::AtomicUsize;
5131
5132 let incarnation = [0x42; CLIENT_INCARNATION_SIZE];
5133 let first = BulkRecord {
5134 id: 1,
5135 kind: BulkKind::Filesystem,
5136 flow: BulkFlow::HostToGuest,
5137 offset: 0,
5138 payload: Bytes::from(vec![0xa5; 512]),
5139 };
5140 let second = BulkRecord {
5141 offset: first.payload.len() as u64,
5142 payload: Bytes::from(vec![0x5a; 512]),
5143 ..first.clone()
5144 };
5145 let mut wire = Vec::new();
5146 for record in [&first, &second] {
5147 wire.extend_from_slice(&incarnation);
5148 codec::encode_bulk_to_buf(record, &mut wire).unwrap();
5149 }
5150 let (mut source, input) = UnixStream::pair().unwrap();
5151 input.set_nonblocking(true).unwrap();
5152 source.write_all(&wire).unwrap();
5153 let mut input = BulkInputState::new(File::from(OwnedFd::from(input))).unwrap();
5154 input.read_buf.truncate(1);
5156
5157 let (mut control_source, control) = UnixStream::pair().unwrap();
5158 control.set_nonblocking(true).unwrap();
5159 let control = AsyncFd::new(control).unwrap();
5160 let (pipe_reader, pipe_writer) = nix::unistd::pipe2(nix::fcntl::OFlag::O_NONBLOCK).unwrap();
5161 let pipe_writer = AsyncFd::new(pipe_writer).unwrap();
5162 let mut writable = pipe_writer.writable().await.unwrap();
5164 loop {
5165 match writable.try_io(|fd| write_to_fd(fd.get_ref().as_raw_fd(), &[0; 4096])) {
5166 Ok(Ok(count)) => assert!(count > 0),
5167 Ok(Err(error)) => panic!("fill stdin pipe: {error}"),
5168 Err(_) => break,
5169 }
5170 }
5171 drop(writable);
5172 let mut drained = [0; 4096];
5173 loop {
5174 match read_from_fd(pipe_reader.as_raw_fd(), &mut drained) {
5175 Ok(count) => assert!(count > 0),
5176 Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => break,
5177 Err(error) => panic!("drain stdin pipe: {error}"),
5178 }
5179 }
5180 control_source.write_all(b"c").unwrap();
5181 let serviced = Arc::new(AtomicUsize::new(0));
5182 let control_serviced = Arc::clone(&serviced);
5183 let control_task = tokio::spawn(async move {
5184 loop {
5185 let mut ready = control.readable().await.unwrap();
5186 let mut byte = [0];
5187 match ready.try_io(|fd| read_from_fd(fd.get_ref().as_raw_fd(), &mut byte)) {
5188 Ok(Ok(1)) => {
5189 assert_eq!(byte, *b"c");
5190 control_serviced.fetch_or(1, Ordering::Relaxed);
5191 break;
5192 }
5193 Ok(result) => panic!("read control marker: {result:?}"),
5194 Err(_) => continue,
5195 }
5196 }
5197 });
5198 let stdin_serviced = Arc::clone(&serviced);
5199 let stdin_task = tokio::spawn(async move {
5200 std::future::poll_fn(|cx| {
5201 loop {
5202 let mut ready = std::task::ready!(pipe_writer.poll_write_ready(cx)).unwrap();
5203 match ready.try_io(|fd| write_to_fd(fd.get_ref().as_raw_fd(), b"i")) {
5204 Ok(result) => return Poll::Ready(result),
5205 Err(_) => continue,
5206 }
5207 }
5208 })
5209 .await
5210 .unwrap();
5211 stdin_serviced.fetch_or(2, Ordering::Relaxed);
5212 });
5213
5214 let mut budget = AgentReadBudget::default();
5215 let mut received = Vec::new();
5216 let mut serviced_before_first_record = false;
5217 while received.len() < 2 {
5218 let frames = tokio::select! {
5219 frames = input.read_turn(&mut budget) => frames.unwrap(),
5220 _ = std::future::pending::<()>() => unreachable!(),
5221 };
5222 for frame in frames {
5223 assert_eq!(frame.incarnation, incarnation);
5224 received.push(frame.record);
5225 }
5226 if yield_at_boundary {
5229 budget.yield_if_exhausted().await;
5230 }
5231 if received.is_empty() && serviced.load(Ordering::Relaxed) == 3 {
5232 serviced_before_first_record = true;
5233 }
5234 }
5235 assert_eq!(
5236 received,
5237 [first, second],
5238 "yield preserves record bytes and FIFO"
5239 );
5240 assert!(input.input.is_empty());
5241 control_task.await.unwrap();
5242 stdin_task.await.unwrap();
5243 let mut marker = [0];
5244 assert_eq!(
5245 read_from_fd(pipe_reader.as_raw_fd(), &mut marker).unwrap(),
5246 1
5247 );
5248 assert_eq!(marker, *b"i");
5249 serviced_before_first_record
5250 }
5251
5252 #[tokio::test(flavor = "current_thread")]
5253 async fn cancelled_bulk_read_keeps_partial_record_and_read_budget() {
5254 use std::io::Write;
5255 use std::os::fd::OwnedFd;
5256 use std::os::unix::net::UnixStream;
5257
5258 let incarnation = [0x29; CLIENT_INCARNATION_SIZE];
5259 let record = BulkRecord {
5260 id: 1,
5261 kind: BulkKind::Filesystem,
5262 flow: BulkFlow::HostToGuest,
5263 offset: 0,
5264 payload: Bytes::from(vec![0x51; 512]),
5265 };
5266 let mut wire = incarnation.to_vec();
5267 codec::encode_bulk_to_buf(&record, &mut wire).unwrap();
5268 let (mut source, input) = UnixStream::pair().unwrap();
5269 input.set_nonblocking(true).unwrap();
5270 let mut input = BulkInputState::new(File::from(OwnedFd::from(input))).unwrap();
5271 let mut budget = AgentReadBudget::default();
5272 source.write_all(&wire[..40]).unwrap();
5273 assert!(input.read_turn(&mut budget).await.unwrap().is_empty());
5274 assert!(input.read_turn(&mut budget).await.unwrap().is_empty());
5276 let calls = budget.calls;
5277 tokio::select! {
5278 biased;
5279 result = input.read_turn(&mut budget) => panic!("unexpected read: {result:?}"),
5280 _ = std::future::ready(()) => {},
5281 }
5282 assert_eq!(input.input.as_ref(), &wire[..40]);
5283 assert_eq!(budget.bytes, 40);
5284 assert_eq!(budget.calls, calls);
5285 source.write_all(&wire[40..]).unwrap();
5286 let frames = input.read_turn(&mut budget).await.unwrap();
5287 assert_eq!(frames.len(), 1);
5288 assert_eq!(frames[0].record, record);
5289 assert_eq!(frames[0].incarnation, incarnation);
5290 assert!(input.input.is_empty());
5291 assert_eq!(budget.bytes, wire.len());
5292 }
5293
5294 #[test]
5295 fn disconnect_cleanup_keeps_other_clients_bulk_offsets() {
5296 let mut state = AgentState::default();
5297 for id in [1, 10, 20] {
5298 state.bulk_received_offsets.insert(id, id as u64);
5299 state.pending_bulk_finishes.insert(
5300 id,
5301 BulkFinish {
5302 kind: BulkKind::Filesystem,
5303 flow: BulkFlow::HostToGuest,
5304 final_offset: id as u64,
5305 },
5306 );
5307 }
5308
5309 clear_bulk_receive_range(&mut state, 10, 20);
5310
5311 assert_eq!(state.bulk_received_offsets.len(), 2);
5312 assert!(state.bulk_received_offsets.contains_key(&1));
5313 assert!(state.bulk_received_offsets.contains_key(&20));
5314 assert_eq!(state.pending_bulk_finishes.len(), 2);
5315 assert!(state.pending_bulk_finishes.contains_key(&1));
5316 assert!(state.pending_bulk_finishes.contains_key(&20));
5317 }
5318
5319 #[test]
5320 fn recycled_range_rejects_old_incarnation_and_stale_disconnect() {
5321 let id_start = 1;
5322 let id_end_exclusive = microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP;
5323 let old = [0x11; CLIENT_INCARNATION_SIZE];
5324 let new = [0x22; CLIENT_INCARNATION_SIZE];
5325 let mut state = AgentState::default();
5326
5327 establish_relay_client(
5328 &mut state,
5329 RelayClientConnected {
5330 id_start,
5331 id_end_exclusive,
5332 incarnation: old,
5333 },
5334 )
5335 .unwrap();
5336 state.bulk_received_offsets.insert(id_start, 4096);
5337 let (session_tx, _session_rx) = SessionOutputSender::channel();
5338 let (removed, cleanup) = disconnect_relay_client(
5339 &mut state,
5340 RelayClientDisconnected {
5341 id_start,
5342 id_end_exclusive,
5343 incarnation: Some(old),
5344 },
5345 &session_tx,
5346 )
5347 .unwrap();
5348 assert!(removed);
5349 assert!(cleanup.is_none());
5350 establish_relay_client(
5351 &mut state,
5352 RelayClientConnected {
5353 id_start,
5354 id_end_exclusive,
5355 incarnation: new,
5356 },
5357 )
5358 .unwrap();
5359
5360 assert_eq!(client_incarnation_for_id(&state, id_start), Some(new));
5361 assert!(validate_bulk_client_incarnation(&state, id_start, new).unwrap());
5362 assert!(!validate_bulk_client_incarnation(&state, id_start, old).unwrap());
5363
5364 let other_start = microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP + 1;
5365 let other_end = microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP * 2;
5366 establish_relay_client(
5367 &mut state,
5368 RelayClientConnected {
5369 id_start: other_start,
5370 id_end_exclusive: other_end,
5371 incarnation: [0x33; CLIENT_INCARNATION_SIZE],
5372 },
5373 )
5374 .unwrap();
5375 assert!(validate_bulk_client_incarnation(&state, other_start, new).is_err());
5376 assert!(!state.bulk_received_offsets.contains_key(&id_start));
5377 assert!(
5378 !disconnect_relay_client(
5379 &mut state,
5380 RelayClientDisconnected {
5381 id_start,
5382 id_end_exclusive,
5383 incarnation: Some(old),
5384 },
5385 &session_tx,
5386 )
5387 .unwrap()
5388 .0
5389 );
5390 assert_eq!(client_incarnation_for_id(&state, id_start), Some(new));
5391 }
5392
5393 #[test]
5394 fn aggregate_bulk_input_budget_is_bounded_and_recoverable() {
5395 let state = AgentState::default();
5396 let budget = state.bulk_input_budget;
5397 let held = Arc::clone(&budget)
5398 .try_acquire_many_owned(BULK_INPUT_BYTE_CAPACITY as u32)
5399 .unwrap();
5400
5401 assert!(Arc::clone(&budget).try_acquire_owned().is_err());
5402 drop(held);
5403 assert_eq!(budget.available_permits(), BULK_INPUT_BYTE_CAPACITY);
5404 }
5405
5406 #[test]
5407 fn disconnect_ack_does_not_refresh_idle_activity() {
5408 let ack = encode_relay_client_disconnected_ack(RelayClientDisconnectedAck {
5409 id_start: 1,
5410 id_end_exclusive: microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP,
5411 incarnation: [0x4a; CLIENT_INCARNATION_SIZE],
5412 });
5413 assert!(!encoded_guest_message_refreshes_idle_timer(
5414 &ack,
5415 0,
5416 ack.len() - 4,
5417 ));
5418 }
5419
5420 #[tokio::test]
5421 async fn closed_optional_receiver_disables_itself_without_spinning() {
5422 let (tx, rx) = tokio::sync::mpsc::channel::<u8>(1);
5423 drop(tx);
5424 let mut receiver = Some(rx);
5425
5426 assert!(
5427 tokio::time::timeout(Duration::from_millis(10), recv_optional(&mut receiver))
5428 .await
5429 .is_err()
5430 );
5431 assert!(receiver.is_none());
5432 }
5433
5434 #[tokio::test]
5435 async fn bulk_output_tombstone_discards_late_producer_record() {
5436 let incarnation = [0x5a; CLIENT_INCARNATION_SIZE];
5437 let (session_tx, _control_rx, mut bulk_rx, _command_rx) =
5438 SessionOutputSender::split_channel();
5439 let owner_tx = session_tx.with_incarnation(Some(incarnation));
5440 let record = || BulkRecord {
5441 id: 17,
5442 kind: BulkKind::Filesystem,
5443 flow: BulkFlow::GuestToHost,
5444 offset: 0,
5445 payload: bytes::Bytes::from_static(b"late"),
5446 };
5447 assert!(
5448 owner_tx
5449 .send(
5450 17,
5451 SessionOutput::Bulk(crate::session::BulkSessionOutput::new(
5452 record(),
5453 RawActivity::default(),
5454 )),
5455 )
5456 .await
5457 );
5458
5459 let mut generation = 0;
5460 let mut flows = HashMap::new();
5461 let mut active = VecDeque::new();
5462 let mut retired = HashMap::new();
5463 let mut retiring_incarnations = HashSet::new();
5464 enqueue_bulk_output(
5465 bulk_rx.recv().await.unwrap(),
5466 0,
5467 &mut flows,
5468 &mut active,
5469 &retired,
5470 &retiring_incarnations,
5471 )
5472 .unwrap();
5473 let (completion, mut completed) = tokio::sync::oneshot::channel();
5474 let cleanup = apply_bulk_output_command(
5475 BulkOutputCommand::DropFlow {
5476 incarnation,
5477 id: 17,
5478 completion,
5479 },
5480 &mut generation,
5481 &mut flows,
5482 &mut active,
5483 &mut retired,
5484 &mut retiring_incarnations,
5485 &mut BulkOutputPosition::default(),
5486 )
5487 .unwrap();
5488 assert!(matches!(
5489 completed.try_recv(),
5490 Err(tokio::sync::oneshot::error::TryRecvError::Empty)
5491 ));
5492 complete_bulk_output_cleanups(vec![cleanup], &mut retired, &mut retiring_incarnations);
5493 assert_eq!(completed.try_recv(), Ok(()));
5494
5495 assert!(
5496 owner_tx
5497 .send(
5498 17,
5499 SessionOutput::Bulk(crate::session::BulkSessionOutput::new(
5500 record(),
5501 RawActivity::default(),
5502 )),
5503 )
5504 .await
5505 );
5506 enqueue_bulk_output(
5507 bulk_rx.recv().await.unwrap(),
5508 0,
5509 &mut flows,
5510 &mut active,
5511 &retired,
5512 &retiring_incarnations,
5513 )
5514 .unwrap();
5515
5516 assert!(flows.is_empty());
5517 assert!(active.is_empty());
5518 assert!(bulk_output_is_retired(&retired, incarnation, 17));
5519 }
5520
5521 #[test]
5522 fn record_encoded_guest_messages_counts_only_appended_frames() {
5523 let mut out_buf = Vec::new();
5524 let existing =
5525 Message::with_payload(MessageType::ExecStarted, 1, &ExecStarted { pid: 123 }).unwrap();
5526 codec::encode_to_buf(&existing, &mut out_buf).unwrap();
5527 let start = out_buf.len();
5528
5529 let appended =
5530 Message::with_payload(MessageType::ExecStarted, 2, &ExecStarted { pid: 456 }).unwrap();
5531 codec::encode_to_buf(&appended, &mut out_buf).unwrap();
5532
5533 let mut activity = ActivityTracker::new();
5534 record_encoded_guest_messages(&out_buf, start, &mut activity);
5535
5536 assert_eq!(activity.activity_seq, 1);
5537 assert_eq!(activity.counters.guest_messages, 1);
5538 }
5539
5540 #[test]
5541 fn apply_raw_activity_updates_guest_and_byte_counters() {
5542 let mut activity = ActivityTracker::new();
5543
5544 apply_raw_activity(RawActivity::fs_bytes(42), &mut activity);
5545 apply_raw_activity(RawActivity::tcp_bytes(7), &mut activity);
5546
5547 assert_eq!(activity.activity_seq, 2);
5548 assert_eq!(activity.counters.guest_messages, 2);
5549 assert_eq!(activity.counters.fs_bytes, 42);
5550 assert_eq!(activity.counters.tcp_bytes, 7);
5551 }
5552
5553 #[test]
5554 fn apply_raw_activity_preserves_coalesced_message_count() {
5555 let mut activity = ActivityTracker::new();
5556
5557 apply_raw_activity(
5558 RawActivity {
5559 guest_messages: 16,
5560 fs_bytes: 4 * 1024 * 1024,
5561 tcp_bytes: 0,
5562 },
5563 &mut activity,
5564 );
5565
5566 assert_eq!(activity.activity_seq, 16);
5567 assert_eq!(activity.counters.guest_messages, 16);
5568 assert_eq!(activity.counters.fs_bytes, 4 * 1024 * 1024);
5569 }
5570
5571 #[tokio::test]
5572 async fn filesystem_finish_remains_admissible_after_a_full_data_window() {
5573 let (records, _record_rx) = tokio::sync::mpsc::channel(FS_BULK_INPUT_ITEM_CAPACITY);
5574 let (finish, mut finish_rx) = tokio::sync::mpsc::channel(1);
5575 let worker = FsBulkWriteWorker {
5576 records,
5577 finish,
5578 task: tokio::spawn(std::future::pending::<()>()),
5579 };
5580 let payload = Bytes::from(vec![0u8; MIN_BULK_RECORD_PAYLOAD as usize]);
5581
5582 for index in 0..FS_BULK_INPUT_ITEM_CAPACITY {
5583 worker
5584 .records
5585 .try_send(AdmittedBulkRecord::for_test(BulkRecord {
5586 id: 17,
5587 kind: BulkKind::Filesystem,
5588 flow: BulkFlow::HostToGuest,
5589 offset: (index * MIN_BULK_RECORD_PAYLOAD as usize) as u64,
5590 payload: payload.clone(),
5591 }))
5592 .unwrap();
5593 }
5594 assert_eq!(worker.records.capacity(), 0);
5595
5596 let expected = BulkFinish {
5597 kind: BulkKind::Filesystem,
5598 flow: BulkFlow::HostToGuest,
5599 final_offset: microsandbox_protocol::bulk::DEFAULT_BULK_WINDOW,
5600 };
5601 worker.finish.try_send(expected).unwrap();
5602 assert_eq!(finish_rx.recv().await, Some(expected));
5603 worker.task.abort();
5604 }
5605
5606 #[test]
5607 fn maintenance_messages_do_not_implicitly_refresh_idle_timer() {
5608 assert!(!message_refreshes_idle_timer(&MessageType::ClockSync));
5609 assert!(!message_refreshes_idle_timer(&MessageType::Ping));
5610 assert!(!message_refreshes_idle_timer(&MessageType::Touch));
5611 assert!(message_refreshes_idle_timer(&MessageType::ExecRequest));
5612 }
5613
5614 #[test]
5615 fn maintenance_replies_do_not_refresh_idle_timer() {
5616 assert!(!guest_message_refreshes_idle_timer(&MessageType::Pong));
5617 assert!(!guest_message_refreshes_idle_timer(&MessageType::Touched));
5618 assert!(!guest_message_refreshes_idle_timer(&MessageType::CoreError));
5619 assert!(guest_message_refreshes_idle_timer(&MessageType::ExecStdout));
5620 }
5621
5622 #[test]
5623 fn record_encoded_guest_messages_ignores_pong_and_touched() {
5624 let mut out_buf = Vec::new();
5625 let pong = Message::with_payload(MessageType::Pong, 1, &Pong {}).unwrap();
5626 codec::encode_to_buf(&pong, &mut out_buf).unwrap();
5627
5628 let touched =
5629 Message::with_payload(MessageType::Touched, 2, &Touched { activity_seq: 42 }).unwrap();
5630 codec::encode_to_buf(&touched, &mut out_buf).unwrap();
5631
5632 let mut activity = ActivityTracker::new();
5633 record_encoded_guest_messages(&out_buf, 0, &mut activity);
5634
5635 assert_eq!(activity.activity_seq, 0);
5636 assert_eq!(activity.counters.guest_messages, 0);
5637 }
5638
5639 #[test]
5640 fn bulk_cancellation_uses_each_operations_existing_terminal_type() {
5641 for (kind, expected) in [
5642 (BulkKind::Filesystem, MessageType::FsResponse),
5643 (BulkKind::Tcp, MessageType::TcpFailed),
5644 ] {
5645 let mut encoded = Vec::new();
5646 encode_bulk_terminal_failure(17, kind, "cancelled".into(), &mut encoded).unwrap();
5647 let mut bytes = BytesMut::from(encoded.as_slice());
5648 let frame = codec::try_decode_frame_from_bytes(&mut bytes)
5649 .unwrap()
5650 .expect("terminal frame");
5651 let DecodedFrame::Control(message) = frame else {
5652 panic!("cancellation terminal must be control");
5653 };
5654 assert_eq!(message.id, 17);
5655 assert_eq!(message.t, expected);
5656 assert_eq!(message.flags, microsandbox_protocol::message::FLAG_TERMINAL);
5657 assert!(bytes.is_empty());
5658 }
5659 }
5660 #[tokio::test]
5661 async fn restore_thaw_checks_ownership_is_idempotent_and_leaves_source_clients_alone() {
5662 use microsandbox_protocol::core::WorkloadThawMode;
5663
5664 let mut state = AgentState::default();
5665 let (mut sender, _output) = SessionOutputSender::channel();
5666 let mut activity = ActivityTracker::new();
5667 let heartbeat = heartbeat::HeartbeatControl::default();
5668 let config = AgentdConfig {
5669 user: None,
5670 security_profile: Default::default(),
5671 default_cwd: None,
5672 default_env: Vec::new(),
5673 };
5674 let mut workload = crate::workload::tests::fake_latch();
5675 let owner = [0x51; CLIENT_INCARNATION_SIZE];
5676 let connect = |state: &mut AgentState| {
5677 establish_relay_client(
5678 state,
5679 RelayClientConnected {
5680 id_start: 1,
5681 id_end_exclusive: microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP,
5682 incarnation: owner,
5683 },
5684 )
5685 .unwrap();
5686 };
5687 connect(&mut state);
5688 workload.freeze("capture").unwrap();
5689
5690 for (attempt, mode, expected) in [
5692 ("wrong", WorkloadThawMode::Restore, MessageType::CoreError),
5693 (
5694 "capture",
5695 WorkloadThawMode::Continue,
5696 MessageType::WorkloadThawed,
5697 ),
5698 ] {
5699 let request = Message::with_payload(
5700 MessageType::WorkloadThaw,
5701 0,
5702 &WorkloadThaw {
5703 attempt_id: attempt.into(),
5704 mode,
5705 },
5706 )
5707 .unwrap();
5708 let mut encoded = Vec::new();
5709 handle_message(
5710 request,
5711 &mut state,
5712 &mut activity,
5713 &mut sender,
5714 &mut encoded,
5715 &config,
5716 &mut workload,
5717 &heartbeat,
5718 )
5719 .await
5720 .unwrap();
5721 let mut bytes = BytesMut::from(encoded.as_slice());
5722 let reply = decode_reply_skipping_credit(&mut bytes);
5723 assert_eq!(reply.t, expected);
5724 assert_eq!(sender.generation(), 0);
5725 assert_eq!(client_incarnation_for_id(&state, 1), Some(owner));
5726 }
5727
5728 for generation in 1..=2 {
5729 let request = Message::with_payload(
5731 MessageType::WorkloadFreeze,
5732 0,
5733 &WorkloadFreeze {
5734 external_mount_tags: Vec::new(),
5735 attempt_id: "capture".into(),
5736 host_input: WorkloadTransportPosition::default(),
5737 },
5738 )
5739 .unwrap();
5740 handle_message(
5741 request,
5742 &mut state,
5743 &mut activity,
5744 &mut sender,
5745 &mut Vec::new(),
5746 &config,
5747 &mut workload,
5748 &heartbeat,
5749 )
5750 .await
5751 .unwrap();
5752 for retry in [false, true] {
5753 let request = Message::with_payload(
5754 MessageType::WorkloadThaw,
5755 0,
5756 &WorkloadThaw {
5757 attempt_id: "capture".into(),
5758 mode: WorkloadThawMode::Restore,
5759 },
5760 )
5761 .unwrap();
5762 let mut encoded = Vec::new();
5763 handle_message(
5764 request,
5765 &mut state,
5766 &mut activity,
5767 &mut sender,
5768 &mut encoded,
5769 &config,
5770 &mut workload,
5771 &heartbeat,
5772 )
5773 .await
5774 .unwrap();
5775 let mut bytes = BytesMut::from(encoded.as_slice());
5776 let reply = decode_reply_skipping_credit(&mut bytes);
5777 assert_eq!(reply.t, MessageType::WorkloadThawed);
5778 assert!(!workload.is_frozen());
5779 assert_eq!(sender.generation(), generation);
5780 assert_eq!(client_incarnation_for_id(&state, 1), retry.then_some(owner));
5781 if !retry {
5782 connect(&mut state);
5784 }
5785 }
5786 }
5787 }
5788
5789 #[tokio::test]
5790 async fn saturated_stdin_allows_thaw_and_preserves_input_and_eof() {
5791 use microsandbox_protocol::core::WorkloadThawMode::{Continue, Restore};
5792
5793 struct StoppedProcessGuard(i32);
5797 impl Drop for StoppedProcessGuard {
5798 fn drop(&mut self) {
5799 unsafe {
5800 libc::kill(-self.0, libc::SIGCONT);
5801 libc::kill(-self.0, libc::SIGKILL);
5802 }
5803 }
5804 }
5805
5806 for (tty, mode, accepted_eof) in [
5807 (false, Continue, false),
5808 (false, Continue, true),
5809 (false, Restore, false),
5810 (false, Restore, true),
5811 (true, Continue, true),
5812 (true, Restore, true),
5813 ] {
5814 let mut state = AgentState::default();
5815 let (mut sender, mut output) = SessionOutputSender::channel();
5816 let mut activity = ActivityTracker::new();
5817 let heartbeat = heartbeat::HeartbeatControl::default();
5818 let config = AgentdConfig {
5819 user: None,
5820 security_profile: Default::default(),
5821 default_cwd: None,
5822 default_env: Vec::new(),
5823 };
5824 let mut workload = crate::workload::tests::fake_latch();
5825 workload.freeze("stdin-cut").unwrap();
5826 let request = ExecRequest {
5827 cmd: "/bin/sh".into(),
5828 args: vec![
5829 "-c".into(),
5830 if tty {
5831 "stty raw -echo; kill -STOP $$; exec cat"
5832 } else {
5833 "kill -STOP $$; exec cat"
5834 }
5835 .into(),
5836 ],
5837 env: vec![],
5838 cwd: None,
5839 user: None,
5840 tty,
5841 rows: 24,
5842 cols: 80,
5843 rlimits: vec![],
5844 };
5845 let session = ExecSession::spawn(
5846 1,
5847 &request,
5848 sender.clone(),
5849 None,
5850 crate::config::SecurityProfile::Default,
5851 None,
5852 )
5853 .unwrap();
5854 let pid = session.pid() as i32;
5855 state.sessions.insert(1, session);
5856 let _cleanup = StoppedProcessGuard(pid);
5857 time::timeout(Duration::from_secs(5), async {
5858 loop {
5859 let status = std::fs::read_to_string(format!("/proc/{pid}/status")).unwrap();
5860 if status.lines().any(|line| line.starts_with("State:\tT")) {
5861 break;
5862 }
5863 time::sleep(Duration::from_millis(1)).await;
5864 }
5865 })
5866 .await
5867 .expect("workload should stop before stdin is admitted");
5868
5869 let mut encoded = Vec::new();
5870 let ledger = state.input_window.clone();
5871 let initial = ledger.credit().unwrap();
5872 let data = vec![0x61; 1024 * 1024];
5873 let charge = ledger.admit(InputLane::Bulk, data.len() + 32).unwrap();
5874 let mut stdin =
5875 Message::with_payload(MessageType::ExecStdin, 1, &ExecStdin { data }).unwrap();
5876 stdin.v = 8;
5879 time::timeout(
5880 Duration::from_millis(100),
5881 handle_message_with_charge(
5882 stdin,
5883 &mut state,
5884 &mut activity,
5885 &mut sender,
5886 &mut encoded,
5887 &config,
5888 &mut workload,
5889 &heartbeat,
5890 Some(charge),
5891 ),
5892 )
5893 .await
5894 .expect("stdin dispatch waited on its blocked consumer")
5895 .unwrap();
5896 assert!(encoded.is_empty());
5897 assert!(state.sessions[&1].has_pending_stdin());
5898 if accepted_eof {
5899 let charge = ledger.admit(InputLane::Bulk, 32).unwrap();
5900 let mut eof = Message::with_payload(
5901 MessageType::ExecStdin,
5902 1,
5903 &ExecStdin { data: Vec::new() },
5904 )
5905 .unwrap();
5906 eof.v = 8;
5907 time::timeout(
5908 Duration::from_millis(100),
5909 handle_message_with_charge(
5910 eof,
5911 &mut state,
5912 &mut activity,
5913 &mut sender,
5914 &mut encoded,
5915 &config,
5916 &mut workload,
5917 &heartbeat,
5918 Some(charge),
5919 ),
5920 )
5921 .await
5922 .expect("EOF dispatch waited on preceding blocked data")
5923 .unwrap();
5924 }
5925 assert_eq!(
5926 ledger.credit().unwrap(),
5927 initial,
5928 "blocked input refunded before consumption"
5929 );
5930 let position = ledger.position();
5931 let thaw = Message::with_payload(
5932 MessageType::WorkloadThaw,
5933 u32::MAX,
5934 &WorkloadThaw {
5935 attempt_id: "stdin-cut".into(),
5936 mode,
5937 },
5938 )
5939 .unwrap();
5940 time::timeout(
5941 Duration::from_millis(100),
5942 handle_message(
5943 thaw,
5944 &mut state,
5945 &mut activity,
5946 &mut sender,
5947 &mut encoded,
5948 &config,
5949 &mut workload,
5950 &heartbeat,
5951 ),
5952 )
5953 .await
5954 .expect("thaw waited on saturated stdin")
5955 .unwrap();
5956 assert!(!workload.is_frozen());
5957 let mut bytes = BytesMut::from(encoded.as_slice());
5958 let reply = decode_reply_skipping_credit(&mut bytes);
5959 assert_eq!(reply.t, MessageType::WorkloadThawed);
5960 assert_eq!(
5961 ledger.position(),
5962 position,
5963 "restore reset cumulative input position"
5964 );
5965 assert_eq!(
5966 ledger.credit().unwrap(),
5967 initial,
5968 "restore discarded accepted input"
5969 );
5970 assert_eq!(unsafe { libc::kill(-pid, libc::SIGCONT) }, 0);
5971 let mut received = Vec::new();
5972 let mut exited = false;
5973 time::timeout(Duration::from_secs(5), async {
5974 while received.len() < 1024 * 1024
5975 || state.sessions.values().chain(state.detached_sessions.values()).any(ExecSession::has_pending_stdin)
5976 || (!tty && (accepted_eof || mode == Restore) && !exited) {
5977 tokio::select! {
5978 (_, _, result) = std::future::poll_fn(|cx| poll_pending_stdin(&mut state, cx)) => result.unwrap(),
5979 envelope = output.recv() => match envelope.unwrap().output {
5980 SessionOutput::Stdout(data) => received.extend(data),
5981 SessionOutput::Exited(code) => { assert_eq!(code, 0); exited = true; },
5982 _ => {},
5983 },
5984 }
5985 }
5986 })
5987 .await
5988 .expect("accepted stdin or ordered EOF did not drain");
5989 assert_eq!(received, vec![0x61; 1024 * 1024]);
5990 assert_eq!(
5991 ledger.credit().unwrap().bulk_bytes,
5992 initial.bulk_bytes + position.bulk_bytes
5993 );
5994 if !tty && mode == Continue && !accepted_eof {
5995 assert!(!exited);
5997 state
5998 .sessions
5999 .get_mut(&1)
6000 .unwrap()
6001 .enqueue_stdin(b"tail".to_vec(), None)
6002 .unwrap();
6003 let envelope = time::timeout(Duration::from_secs(5), output.recv())
6004 .await
6005 .unwrap()
6006 .unwrap();
6007 assert!(
6008 matches!(envelope.output, SessionOutput::Stdout(ref data) if data == b"tail")
6009 );
6010 }
6011 }
6012 }
6013
6014 #[tokio::test]
6015 async fn restore_detaches_piped_and_pty_workloads_without_killing_or_reusing_output() {
6016 for tty in [false, true] {
6017 let mut state = AgentState::default();
6018 let (mut sender, mut output) = SessionOutputSender::channel();
6019 let id = 1;
6020 let end = microsandbox_protocol::AGENT_RELAY_ID_RANGE_STEP;
6021 let old_owner = [0x31; CLIENT_INCARNATION_SIZE];
6022 let new_owner = [0x32; CLIENT_INCARNATION_SIZE];
6023 establish_relay_client(
6024 &mut state,
6025 RelayClientConnected {
6026 id_start: id,
6027 id_end_exclusive: end,
6028 incarnation: old_owner,
6029 },
6030 )
6031 .unwrap();
6032 let request = |script: &str| ExecRequest {
6033 cmd: "/bin/sh".into(),
6034 args: vec!["-c".into(), script.into()],
6035 env: vec![],
6036 cwd: None,
6037 user: None,
6038 tty,
6039 rows: 24,
6040 cols: 80,
6041 rlimits: vec![],
6042 };
6043 let session = ExecSession::spawn(
6044 id,
6045 &request("echo before; sleep 0.3; echo inherited; exit 23"),
6046 sender.with_incarnation(Some(old_owner)),
6047 None,
6048 crate::config::SecurityProfile::Default,
6049 None,
6050 )
6051 .unwrap();
6052 let pid = session.pid();
6053 state.sessions.insert(id, session);
6054 let before = time::timeout(Duration::from_secs(5), output.recv())
6055 .await
6056 .unwrap()
6057 .unwrap();
6058 assert!(matches!(before.output, SessionOutput::Stdout(_)));
6059 let old_generation = sender.generation();
6060
6061 restore_client_state(&mut state, &mut sender).await.unwrap();
6062 assert_eq!(
6063 unsafe { libc::kill(pid as i32, 0) },
6064 0,
6065 "restore killed the captured process"
6066 );
6067 assert_eq!(state.detached_sessions.len(), 1);
6068 assert!(state.sessions.is_empty());
6069 establish_relay_client(
6070 &mut state,
6071 RelayClientConnected {
6072 id_start: id,
6073 id_end_exclusive: end,
6074 incarnation: new_owner,
6075 },
6076 )
6077 .unwrap();
6078 let fresh = ExecSession::spawn(
6079 id,
6080 &request("echo fresh; sleep 0.8; exit 24"),
6081 sender.with_incarnation(Some(new_owner)),
6082 None,
6083 crate::config::SecurityProfile::Default,
6084 None,
6085 )
6086 .unwrap();
6087 state.sessions.insert(id, fresh);
6088 let mut saw_old_exit = false;
6089 let mut saw_new_exit = false;
6090 let mut saw_inherited = false;
6091 time::timeout(Duration::from_secs(5), async {
6092 while !saw_old_exit || !saw_new_exit {
6093 let envelope = output.recv().await.unwrap();
6094 if envelope.generation == old_generation {
6095 if let SessionOutput::Stdout(bytes) = &envelope.output {
6096 saw_inherited |= String::from_utf8_lossy(bytes).contains("inherited");
6097 }
6098 if let SessionOutput::Exited(code) = &envelope.output {
6099 assert_eq!(*code, 23);
6100 saw_old_exit = true;
6101 }
6102 assert!(discard_inherited_output(
6103 &mut state,
6104 &envelope,
6105 sender.generation()
6106 ));
6107 assert!(
6108 state.sessions.contains_key(&id),
6109 "old exit removed new session"
6110 );
6111 } else {
6112 assert!(!discard_inherited_output(
6113 &mut state,
6114 &envelope,
6115 sender.generation()
6116 ));
6117 if let SessionOutput::Exited(code) = envelope.output {
6118 assert_eq!(code, 24);
6119 saw_new_exit = true;
6120 }
6121 }
6122 }
6123 })
6124 .await
6125 .unwrap();
6126 assert!(saw_inherited, "inherited output reader stopped draining");
6127 assert!(
6128 state.detached_sessions.is_empty(),
6129 "exited inherited handles leaked"
6130 );
6131 }
6132 }
6133
6134 #[tokio::test]
6135 async fn restore_closes_inherited_filesystem_handles_and_tcp_transfers() {
6136 use microsandbox_protocol::fs::{FsOp, FsOpenOptions, FsResponse, FsResponseData};
6137 use tokio::io::AsyncReadExt;
6138
6139 async fn fs_request(
6140 state: &mut AgentState,
6141 sender: &SessionOutputSender,
6142 op: FsOp,
6143 ) -> FsResponse {
6144 let mut encoded = Vec::new();
6145 crate::fs::handle_fs_request(
6146 1,
6147 PROTOCOL_VERSION,
6148 FsRequest { op, bulk: None },
6149 &mut state.fs,
6150 &mut encoded,
6151 sender,
6152 )
6153 .await
6154 .unwrap();
6155 let mut bytes = BytesMut::from(encoded.as_slice());
6156 let Some(DecodedFrame::Control(reply)) =
6157 codec::try_decode_frame_from_bytes(&mut bytes).unwrap()
6158 else {
6159 panic!("missing filesystem response");
6160 };
6161 reply.payload::<FsResponse>().unwrap()
6162 }
6163
6164 let mut state = AgentState::default();
6165 let (mut sender, _output) = SessionOutputSender::channel();
6166 let open = || FsOp::OpenFile {
6167 path: "/dev/null".into(),
6168 options: FsOpenOptions {
6169 read: true,
6170 ..Default::default()
6171 },
6172 };
6173 let Some(FsResponseData::Handle(old_handle)) =
6174 fs_request(&mut state, &sender, open()).await.data
6175 else {
6176 panic!("file did not open");
6177 };
6178 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6179 state.tcp_sessions.insert(
6180 2,
6181 TcpSession::open(
6182 2,
6183 TcpConnect {
6184 host: "127.0.0.1".into(),
6185 port: listener.local_addr().unwrap().port(),
6186 bulk: None,
6187 },
6188 &sender,
6189 ),
6190 );
6191 let (mut peer, _) = time::timeout(Duration::from_secs(5), listener.accept())
6192 .await
6193 .unwrap()
6194 .unwrap();
6195
6196 restore_client_state(&mut state, &mut sender).await.unwrap();
6197 assert!(state.tcp_sessions.is_empty());
6198 assert_eq!(
6199 time::timeout(Duration::from_secs(5), peer.read(&mut [0u8; 1]))
6200 .await
6201 .unwrap()
6202 .unwrap(),
6203 0
6204 );
6205 assert!(
6206 !fs_request(&mut state, &sender, FsOp::FStat { handle: old_handle })
6207 .await
6208 .ok
6209 );
6210 let Some(FsResponseData::Handle(new_handle)) =
6211 fs_request(&mut state, &sender, open()).await.data
6212 else {
6213 panic!("fresh file did not open");
6214 };
6215 assert_ne!(new_handle, old_handle);
6217 assert!(
6218 fs_request(&mut state, &sender, FsOp::FStat { handle: new_handle })
6219 .await
6220 .ok
6221 );
6222 }
6223
6224 #[tokio::test]
6225 async fn restore_retires_unleased_output_and_can_repeat_with_the_same_ids() {
6226 let mut state = AgentState::default();
6227 let (mut sender, mut output) = SessionOutputSender::channel();
6228 for _ in 0..3 {
6229 let old = sender.clone();
6230 old.send(1, SessionOutput::Exited(7)).await;
6231 restore_client_state(&mut state, &mut sender).await.unwrap();
6232 let stale = output.recv().await.unwrap();
6233 assert!(discard_inherited_output(
6234 &mut state,
6235 &stale,
6236 sender.generation()
6237 ));
6238 old.send(1, SessionOutput::Exited(7)).await;
6240 assert!(discard_inherited_output(
6241 &mut state,
6242 &output.recv().await.unwrap(),
6243 sender.generation()
6244 ));
6245 sender.send(1, SessionOutput::Exited(8)).await;
6246 assert!(!discard_inherited_output(
6247 &mut state,
6248 &output.recv().await.unwrap(),
6249 sender.generation()
6250 ));
6251 }
6252 }
6253
6254 #[tokio::test]
6255 async fn restore_bulk_cut_drops_queued_and_late_records_and_releases_capacity() {
6256 let (mut sender, _control, mut bulk, mut commands) = SessionOutputSender::split_channel();
6257 let old = sender.with_incarnation(Some([0x41; CLIENT_INCARNATION_SIZE]));
6258 let send = |sender: SessionOutputSender| async move {
6259 sender
6260 .send(
6261 1,
6262 SessionOutput::Bulk(crate::session::BulkSessionOutput::new(
6263 BulkRecord {
6264 id: 1,
6265 kind: BulkKind::Filesystem,
6266 flow: BulkFlow::GuestToHost,
6267 offset: 0,
6268 payload: bytes::Bytes::from_static(b"old"),
6269 },
6270 RawActivity::default(),
6271 )),
6272 )
6273 .await
6274 };
6275 assert!(send(old.clone()).await);
6276 let reset = tokio::spawn(async move {
6277 sender.restore_generation().await.unwrap();
6278 sender
6279 });
6280 let command = commands.recv().await.unwrap();
6281 let mut generation = 0;
6282 let mut flows = HashMap::new();
6283 let mut active = VecDeque::new();
6284 let mut retired = HashMap::new();
6285 let mut retiring = HashSet::new();
6286 let cleanup = apply_bulk_output_command(
6287 command,
6288 &mut generation,
6289 &mut flows,
6290 &mut active,
6291 &mut retired,
6292 &mut retiring,
6293 &mut BulkOutputPosition::default(),
6294 )
6295 .unwrap();
6296 enqueue_bulk_output(
6297 bulk.recv().await.unwrap(),
6298 generation,
6299 &mut flows,
6300 &mut active,
6301 &retired,
6302 &retiring,
6303 )
6304 .unwrap();
6305 assert!(flows.is_empty());
6306 complete_bulk_output_cleanups(vec![cleanup], &mut retired, &mut retiring);
6307 let fresh = reset.await.unwrap();
6308 assert!(send(old).await);
6309 enqueue_bulk_output(
6310 bulk.recv().await.unwrap(),
6311 generation,
6312 &mut flows,
6313 &mut active,
6314 &retired,
6315 &retiring,
6316 )
6317 .unwrap();
6318 assert!(flows.is_empty());
6319 assert!(send(fresh.with_incarnation(Some([0x42; CLIENT_INCARNATION_SIZE]))).await);
6320 enqueue_bulk_output(
6321 bulk.recv().await.unwrap(),
6322 generation,
6323 &mut flows,
6324 &mut active,
6325 &retired,
6326 &retiring,
6327 )
6328 .unwrap();
6329 assert_eq!(flows.len(), 1);
6330 }
6331
6332 #[test]
6333 fn restore_activity_backpressure_does_not_block_scheduler_control() {
6334 let (sender, mut receiver) = tokio::sync::mpsc::channel(1);
6335 sender.try_send(RawActivity::default()).unwrap();
6336 let mut pending = RawActivity {
6337 guest_messages: 3,
6338 fs_bytes: 17,
6339 ..Default::default()
6340 };
6341 publish_bulk_activity(&sender, &mut pending).unwrap();
6342 assert_eq!(pending.fs_bytes, 17);
6343 receiver.try_recv().unwrap();
6344 publish_bulk_activity(&sender, &mut pending).unwrap();
6345 assert_eq!(receiver.try_recv().unwrap().fs_bytes, 17);
6346 assert_eq!(pending.fs_bytes, 0);
6347 }
6348}