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