arcbox_computer_runtime/agent/
vm_proto.rs

1//! Host-side vsock client for communicating with the in-VM guest agent.
2//!
3//! ## Reaching the guest
4//!
5//! Every protocol function here dials the guest through the driver port's
6//! [`Vsock`] capability: [`connect_to_port`] asks it for a connection to a
7//! guest port and retries — only — while the guest has no listener there
8//! yet, then turns the [`VsockConn`] it hands back into the tokio stream the
9//! frame codec speaks over. Where that connection comes from is the
10//! driver's business.
11//!
12//! The capability comes off the running VM's handle (`HandleVsock` keeps
13//! that handle alive for the callers that hold the vsock across awaits).
14//! In the other direction — the guest agent's readiness dial-out — the
15//! boot pre-listens through the driver's `VsockListen` capability and
16//! [`wait_ready`] takes the one connection.
17//!
18//! ## Frame format
19//!
20//! The opcodes, payload layouts, and the JSON DTOs ([`StartCommand`],
21//! [`WaitPortReq`]) are the exec-channel vocabulary in
22//! [`arcbox_vm_proto::exec`], re-exported here. This module holds the
23//! connection and framing layer; the protocols each have a file of their
24//! own — `exec` (`run`, `exec`), `clock` (`sync_clock`), `net`
25//! (`reconfigure_network`), `wait_port` (`wait_for_port`), `files` — and
26//! `client` is what maps them onto the [`GuestAgent`](crate::agent::GuestAgent)
27//! port. The vocabulary they speak belongs to the port, not here.
28
29use std::time::Duration;
30
31use arcbox_vm_driver::{IoMode, Vsock, VsockConn, VsockListener};
32use tokio::io::{AsyncReadExt, AsyncWriteExt};
33use tokio::net::UnixStream;
34use tokio::sync::mpsc;
35use tracing::warn;
36
37use crate::agent::{ExitStatus, OutputChunk};
38use crate::error::{Result, VmmError};
39
40// Exec-channel vocabulary, shared with vm-agent through `arcbox-vm-proto`.
41pub use arcbox_vm_proto::exec::{AGENT_PORT, MSG_WAIT_PORT, READY_PORT, StartCommand, WaitPortReq};
42pub(crate) use arcbox_vm_proto::exec::{MAX_FRAME_SIZE, MSG_CLOCK_SYNC, MSG_NET_RECONFIG};
43use arcbox_vm_proto::exec::{MSG_EXIT, MSG_STDERR, MSG_STDOUT};
44
45mod client;
46mod clock;
47mod exec;
48mod files;
49mod handle;
50mod net;
51mod wait_port;
52
53pub use client::{VmProtoAgent, VmProtoAgentFactory};
54pub(crate) use handle::HandleVsock;
55
56impl ExitStatus {
57    /// Decode a `MSG_EXIT` payload.
58    ///
59    /// New agents send `[i32 LE code][i32 LE signal]`; agents from before the
60    /// signal extension (e.g. inside restored snapshots) send only the 4-byte
61    /// code, in which case a signal death arrives collapsed as `128 + signal`.
62    fn from_exit_payload(payload: &[u8]) -> Self {
63        if payload.len() >= 8 {
64            let signal = i32::from_le_bytes(payload[4..8].try_into().unwrap());
65            if signal != 0 {
66                return Self::Signaled(signal);
67            }
68        }
69        let code = if payload.len() >= 4 {
70            i32::from_le_bytes(payload[..4].try_into().unwrap())
71        } else {
72            0
73        };
74        Self::Code(code)
75    }
76}
77
78/// How long to wait for the guest agent to start accepting vsock connections.
79const AGENT_READY_TIMEOUT: Duration = Duration::from_secs(30);
80/// First retry delay of the vsock connect backoff; doubles per attempt.
81const AGENT_READY_INITIAL_BACKOFF: Duration = Duration::from_millis(2);
82/// Ceiling for the vsock connect retry backoff.
83const AGENT_READY_MAX_BACKOFF: Duration = Duration::from_millis(200);
84
85/// Next delay of the exponential connect backoff: double, capped at
86/// [`AGENT_READY_MAX_BACKOFF`].
87fn next_backoff(current: Duration) -> Duration {
88    current.saturating_mul(2).min(AGENT_READY_MAX_BACKOFF)
89}
90
91/// Open a host-initiated vsock connection to the guest agent (port 52).
92///
93/// Retries the dial until the guest agent accepts or [`AGENT_READY_TIMEOUT`]
94/// elapses: a guest still booting (kernel up, vm-agent not yet listening)
95/// answers every dial with `ConnectionRefused`, the one transient outcome.
96async fn connect_to_agent(vsock: &dyn Vsock) -> Result<UnixStream> {
97    connect_to_port(vsock, AGENT_PORT).await
98}
99
100/// Open a host-initiated vsock connection to an arbitrary guest port.
101///
102/// Same retry semantics as [`connect_to_agent`].  Used by the file I/O and
103/// port-forward modules which operate on different vsock ports. Retries
104/// only on the two transient answers [`is_transient`] names — no guest
105/// listener yet, or a VM frozen for a checkpoint; every other error is
106/// final.
107pub(crate) async fn connect_to_port(vsock: &dyn Vsock, port: u32) -> Result<UnixStream> {
108    let deadline = tokio::time::Instant::now() + AGENT_READY_TIMEOUT;
109    let mut backoff = AGENT_READY_INITIAL_BACKOFF;
110    loop {
111        match vsock.dial(port).await {
112            Ok(conn) => return into_unix_stream(conn),
113            Err(error) if is_transient(&error) => {}
114            Err(error) => return Err(error.into()),
115        }
116        if tokio::time::Instant::now() >= deadline {
117            return Err(VmmError::Vsock(format!(
118                "vsock port {port} did not become ready within {}s",
119                AGENT_READY_TIMEOUT.as_secs(),
120            )));
121        }
122        tokio::time::sleep(backoff).await;
123        backoff = next_backoff(backoff);
124    }
125}
126
127/// The dial outcomes [`connect_to_port`] retries: the guest has no listener
128/// on the port yet ([`arcbox_vm_driver::Error::Io`] of kind
129/// [`ConnectionRefused`](std::io::ErrorKind::ConnectionRefused), the port's
130/// answer for that), or the VM is frozen for a checkpoint
131/// ([`arcbox_vm_driver::VmState::Quiesced`]). The second is a moment, not
132/// a state to fail on: a file read or an exec that races a checkpoint used
133/// to sit on the VMM's connect until the guest resumed, and the driver's
134/// refusal to dial a quiesced VM must not turn that wait into an error. A
135/// VM that stays frozen — pause holds it, then kills it — turns the retry
136/// into the exited-VM error or the ready budget's.
137fn is_transient(error: &arcbox_vm_driver::Error) -> bool {
138    match error {
139        arcbox_vm_driver::Error::Io(io) => io.kind() == std::io::ErrorKind::ConnectionRefused,
140        arcbox_vm_driver::Error::WrongState { state, .. } => {
141            *state == arcbox_vm_driver::VmState::Quiesced
142        }
143        _ => false,
144    }
145}
146
147/// Register a dialed connection with the tokio reactor.
148///
149/// This client is tokio tasks all the way down and has no blocking-thread
150/// transport, so a connection whose driver requires [`IoMode::Blocking`] is
151/// refused rather than misregistered with the reactor.
152fn into_unix_stream(conn: VsockConn) -> Result<UnixStream> {
153    match conn.mode {
154        IoMode::Async => {
155            let stream = std::os::unix::net::UnixStream::from(conn.fd);
156            stream.set_nonblocking(true)?;
157            Ok(UnixStream::from_std(stream)?)
158        }
159        IoMode::Blocking => Err(VmmError::Vsock(
160            "vsock connection requires blocking I/O, which the guest-agent client cannot drive"
161                .into(),
162        )),
163    }
164}
165
166/// Wait for the guest agent's readiness dial-out on a pre-bound listener:
167/// `accept()` one connection and read (and discard) its single byte.
168/// Completion is the readiness event.
169pub(crate) async fn wait_ready(listener: &mut dyn VsockListener) -> Result<()> {
170    let conn = listener
171        .accept()
172        .await
173        .map_err(|e| VmmError::Vsock(format!("accept on ready socket: {e}")))?;
174    let mut stream = into_unix_stream(conn)?;
175    let mut byte = [0u8; 1];
176    stream
177        .read(&mut byte)
178        .await
179        .map_err(|e| VmmError::Vsock(format!("read ready byte: {e}")))?;
180    Ok(())
181}
182
183/// Write a single frame to any `AsyncWrite`.
184pub(crate) async fn write_frame<W: AsyncWriteExt + Unpin>(
185    w: &mut W,
186    msg_type: u8,
187    payload: &[u8],
188) -> std::io::Result<()> {
189    if payload.len() > MAX_FRAME_SIZE {
190        return Err(std::io::Error::new(
191            std::io::ErrorKind::InvalidInput,
192            format!(
193                "frame payload too large: {} bytes (max {MAX_FRAME_SIZE})",
194                payload.len()
195            ),
196        ));
197    }
198    w.write_u8(msg_type).await?;
199    w.write_u32_le(payload.len() as u32).await?;
200    if !payload.is_empty() {
201        w.write_all(payload).await?;
202    }
203    Ok(())
204}
205
206/// Read a single frame from any `AsyncRead`.
207pub(crate) async fn read_frame<R: AsyncReadExt + Unpin>(
208    r: &mut R,
209) -> std::io::Result<(u8, Vec<u8>)> {
210    let msg_type = r.read_u8().await?;
211    let len = r.read_u32_le().await? as usize;
212    if len > MAX_FRAME_SIZE {
213        return Err(std::io::Error::new(
214            std::io::ErrorKind::InvalidData,
215            format!("frame too large: {len} bytes (max {MAX_FRAME_SIZE})"),
216        ));
217    }
218    let mut payload = vec![0u8; len];
219    if len > 0 {
220        r.read_exact(&mut payload).await?;
221    }
222    Ok((msg_type, payload))
223}
224
225/// Drain an output half, forwarding frames to `tx` until `MSG_EXIT` or error.
226async fn drain_output<R: AsyncReadExt + Unpin>(
227    mut read_half: R,
228    tx: mpsc::Sender<Result<OutputChunk>>,
229) {
230    loop {
231        match read_frame(&mut read_half).await {
232            Ok((msg_type, payload)) => {
233                let chunk = match msg_type {
234                    MSG_STDOUT => OutputChunk::Stdout(payload),
235                    MSG_STDERR => OutputChunk::Stderr(payload),
236                    MSG_EXIT => {
237                        let status = ExitStatus::from_exit_payload(&payload);
238                        let _ = tx.send(Ok(OutputChunk::Exit(status))).await;
239                        break;
240                    }
241                    other => {
242                        warn!(msg_type = other, "unknown agent→host frame type; ignoring");
243                        continue;
244                    }
245                };
246                if tx.send(Ok(chunk)).await.is_err() {
247                    break;
248                }
249            }
250            Err(e) => {
251                let _ = tx
252                    .send(Err(VmmError::Vsock(format!("agent read error: {e}"))))
253                    .await;
254                break;
255            }
256        }
257    }
258}
259
260#[cfg(test)]
261mod tests {
262    use std::os::fd::OwnedFd;
263
264    use arcbox_vm_proto::exec::{MSG_EOF, MSG_RESIZE, MSG_SIGNAL, MSG_START};
265    use async_trait::async_trait;
266
267    use super::*;
268
269    /// Build a raw frame byte-by-byte for use in read tests.
270    fn make_raw_frame(msg_type: u8, payload: &[u8]) -> Vec<u8> {
271        let mut buf = Vec::new();
272        buf.push(msg_type);
273        buf.extend_from_slice(&(payload.len() as u32).to_le_bytes());
274        buf.extend_from_slice(payload);
275        buf
276    }
277
278    #[test]
279    fn connect_backoff_doubles_up_to_the_cap() {
280        let mut delays = Vec::new();
281        let mut backoff = AGENT_READY_INITIAL_BACKOFF;
282        for _ in 0..10 {
283            delays.push(backoff.as_millis());
284            backoff = next_backoff(backoff);
285        }
286        assert_eq!(delays, [2, 4, 8, 16, 32, 64, 128, 200, 200, 200]);
287    }
288
289    /// A [`Vsock`] whose dials answer from a script. Once the script runs
290    /// out it answers `ConnectionRefused` forever — a guest that never
291    /// listens.
292    struct ScriptedVsock {
293        script: std::sync::Mutex<std::collections::VecDeque<arcbox_vm_driver::Result<IoMode>>>,
294        dials: std::sync::atomic::AtomicUsize,
295    }
296
297    impl ScriptedVsock {
298        fn new(script: impl IntoIterator<Item = arcbox_vm_driver::Result<IoMode>>) -> Self {
299            Self {
300                script: std::sync::Mutex::new(script.into_iter().collect()),
301                dials: std::sync::atomic::AtomicUsize::new(0),
302            }
303        }
304
305        fn dials(&self) -> usize {
306            self.dials.load(std::sync::atomic::Ordering::SeqCst)
307        }
308    }
309
310    #[async_trait]
311    impl Vsock for ScriptedVsock {
312        async fn dial(&self, _port: u32) -> arcbox_vm_driver::Result<VsockConn> {
313            self.dials.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
314            let next = self.script.lock().unwrap().pop_front();
315            let mode = next.unwrap_or_else(|| Err(refused()))?;
316            let (ours, _theirs) = std::os::unix::net::UnixStream::pair().unwrap();
317            Ok(VsockConn {
318                fd: OwnedFd::from(ours),
319                mode,
320            })
321        }
322    }
323
324    fn refused() -> arcbox_vm_driver::Error {
325        std::io::Error::from(std::io::ErrorKind::ConnectionRefused).into()
326    }
327
328    fn quiesced() -> arcbox_vm_driver::Error {
329        arcbox_vm_driver::Error::WrongState {
330            id: arcbox_vm_driver::VmId::new("vm").unwrap(),
331            state: arcbox_vm_driver::VmState::Quiesced,
332            expected: "running",
333        }
334    }
335
336    #[tokio::test]
337    async fn connect_retries_only_while_the_guest_is_not_listening() {
338        let vsock = ScriptedVsock::new([Err(refused()), Err(refused()), Ok(IoMode::Async)]);
339        connect_to_port(&vsock, AGENT_PORT).await.unwrap();
340        assert_eq!(vsock.dials(), 3);
341    }
342
343    /// A VM frozen for a checkpoint is dialed again once it thaws — the
344    /// wait a caller racing a checkpoint always had — while an exited VM
345    /// is final.
346    #[tokio::test]
347    async fn connect_waits_out_a_quiesced_vm() {
348        let vsock = ScriptedVsock::new([Err(quiesced()), Err(quiesced()), Ok(IoMode::Async)]);
349        connect_to_port(&vsock, AGENT_PORT).await.unwrap();
350        assert_eq!(vsock.dials(), 3);
351    }
352
353    #[tokio::test]
354    async fn connect_fails_fast_on_any_other_error() {
355        // Each final error keeps its native shape through the conversion:
356        // an I/O failure stays `Io`, a driver `WrongState` stays
357        // `WrongState` (the guest agent maps that to 412, not 500).
358        type Native = fn(&VmmError) -> bool;
359        let finals: [(arcbox_vm_driver::Error, Native); 2] = [
360            (
361                arcbox_vm_driver::Error::Io(std::io::ErrorKind::BrokenPipe.into()),
362                |err| matches!(err, VmmError::Io(_)),
363            ),
364            (
365                arcbox_vm_driver::Error::WrongState {
366                    id: arcbox_vm_driver::VmId::new("vm").unwrap(),
367                    state: arcbox_vm_driver::VmState::Exited(
368                        arcbox_vm_driver::ExitStatus::signaled(9),
369                    ),
370                    expected: "running",
371                },
372                |err| matches!(err, VmmError::WrongState { expected, actual, .. } if expected == "running" && actual.starts_with("exited")),
373            ),
374        ];
375        for (error, native) in finals {
376            let vsock = ScriptedVsock::new([Err(error), Ok(IoMode::Async)]);
377            let err = connect_to_port(&vsock, AGENT_PORT).await.unwrap_err();
378            assert!(native(&err), "unexpected error: {err}");
379            assert_eq!(vsock.dials(), 1, "a final error must not be retried");
380        }
381    }
382
383    #[tokio::test(start_paused = true)]
384    async fn connect_gives_up_after_the_ready_budget() {
385        let vsock = ScriptedVsock::new([]);
386        let err = connect_to_port(&vsock, AGENT_PORT).await.unwrap_err();
387        assert!(
388            matches!(err, VmmError::Vsock(ref m) if m.contains("did not become ready within 30s")),
389            "unexpected error: {err}"
390        );
391        assert!(vsock.dials() > 1);
392    }
393
394    #[tokio::test]
395    async fn blocking_connections_are_refused_not_misregistered() {
396        let vsock = ScriptedVsock::new([Ok(IoMode::Blocking)]);
397        let err = connect_to_port(&vsock, AGENT_PORT).await.unwrap_err();
398        assert!(
399            matches!(err, VmmError::Vsock(ref m) if m.contains("blocking I/O")),
400            "unexpected error: {err}"
401        );
402    }
403
404    #[tokio::test]
405    async fn test_write_read_frame_roundtrip() {
406        let (mut a, mut b) = tokio::io::duplex(256);
407        write_frame(&mut a, MSG_START, b"hello world")
408            .await
409            .unwrap();
410        let (msg_type, payload) = read_frame(&mut b).await.unwrap();
411        assert_eq!(msg_type, MSG_START);
412        assert_eq!(payload, b"hello world");
413    }
414
415    #[tokio::test]
416    async fn test_empty_payload_frame() {
417        let (mut a, mut b) = tokio::io::duplex(64);
418        write_frame(&mut a, MSG_EOF, &[]).await.unwrap();
419        let (msg_type, payload) = read_frame(&mut b).await.unwrap();
420        assert_eq!(msg_type, MSG_EOF);
421        assert!(payload.is_empty());
422    }
423
424    #[tokio::test]
425    async fn test_exit_code_encoding() {
426        let exit_code: i32 = 42;
427        let (mut a, mut b) = tokio::io::duplex(64);
428        write_frame(&mut a, MSG_EXIT, &exit_code.to_le_bytes())
429            .await
430            .unwrap();
431        let (msg_type, payload) = read_frame(&mut b).await.unwrap();
432        assert_eq!(msg_type, MSG_EXIT);
433        let decoded = i32::from_le_bytes(payload[..4].try_into().unwrap());
434        assert_eq!(decoded, 42);
435    }
436
437    #[test]
438    fn exit_payload_decodes_legacy_and_signal_forms() {
439        // Legacy 4-byte form (old vm-agent): always a plain code.
440        assert_eq!(
441            ExitStatus::from_exit_payload(&7i32.to_le_bytes()),
442            ExitStatus::Code(7)
443        );
444        // 8-byte form, signal 0: a normal exit — even for code 137, which the
445        // legacy form could not distinguish from a SIGKILL death.
446        let mut normal_137 = Vec::new();
447        normal_137.extend_from_slice(&137i32.to_le_bytes());
448        normal_137.extend_from_slice(&0i32.to_le_bytes());
449        assert_eq!(
450            ExitStatus::from_exit_payload(&normal_137),
451            ExitStatus::Code(137)
452        );
453        // 8-byte form, signal set: a signal death.
454        let mut sigkill = Vec::new();
455        sigkill.extend_from_slice(&137i32.to_le_bytes());
456        sigkill.extend_from_slice(&9i32.to_le_bytes());
457        assert_eq!(
458            ExitStatus::from_exit_payload(&sigkill),
459            ExitStatus::Signaled(9)
460        );
461        assert_eq!(ExitStatus::Signaled(9).conventional_code(), 137);
462        // Truncated payload degrades to code 0 (matches the old lenient parse).
463        assert_eq!(ExitStatus::from_exit_payload(&[1, 2]), ExitStatus::Code(0));
464    }
465
466    #[tokio::test]
467    async fn signal_frame_round_trips() {
468        let (mut a, mut b) = tokio::io::duplex(64);
469        write_frame(&mut a, MSG_SIGNAL, &15i32.to_le_bytes())
470            .await
471            .unwrap();
472        let (msg_type, payload) = read_frame(&mut b).await.unwrap();
473        assert_eq!(msg_type, MSG_SIGNAL);
474        assert_eq!(i32::from_le_bytes(payload[..4].try_into().unwrap()), 15);
475    }
476
477    #[tokio::test]
478    async fn test_resize_frame_encoding() {
479        let width: u16 = 80;
480        let height: u16 = 24;
481        let mut resize_payload = [0u8; 4];
482        resize_payload[..2].copy_from_slice(&width.to_le_bytes());
483        resize_payload[2..].copy_from_slice(&height.to_le_bytes());
484
485        let (mut a, mut b) = tokio::io::duplex(64);
486        write_frame(&mut a, MSG_RESIZE, &resize_payload)
487            .await
488            .unwrap();
489        let (msg_type, payload) = read_frame(&mut b).await.unwrap();
490        assert_eq!(msg_type, MSG_RESIZE);
491        let w = u16::from_le_bytes(payload[..2].try_into().unwrap());
492        let h = u16::from_le_bytes(payload[2..].try_into().unwrap());
493        assert_eq!(w, 80);
494        assert_eq!(h, 24);
495    }
496
497    #[tokio::test]
498    async fn test_read_frame_from_raw_bytes() {
499        // Verify the parser accepts hand-crafted bytes (regression guard).
500        let raw = make_raw_frame(MSG_STDOUT, b"output line\n");
501        let mut cursor = std::io::Cursor::new(raw);
502        let (msg_type, payload) = read_frame(&mut cursor).await.unwrap();
503        assert_eq!(msg_type, MSG_STDOUT);
504        assert_eq!(payload, b"output line\n");
505    }
506
507    #[test]
508    fn test_start_command_json_serde() {
509        let cmd = StartCommand {
510            cmd: vec!["echo".into(), "hello".into()],
511            env: std::collections::HashMap::new(),
512            working_dir: "/tmp".into(),
513            user: "root".into(),
514            tty: false,
515            tty_width: 0,
516            tty_height: 0,
517            timeout_seconds: 30,
518        };
519        let json = serde_json::to_string(&cmd).unwrap();
520        let decoded: StartCommand = serde_json::from_str(&json).unwrap();
521        assert_eq!(decoded.cmd, vec!["echo", "hello"]);
522        assert_eq!(decoded.working_dir, "/tmp");
523        assert_eq!(decoded.timeout_seconds, 30);
524        assert!(!decoded.tty);
525    }
526}
arcbox_computer_runtime/agent/vm_proto.rs

arcbox_computer_runtime/agent/
vm_proto.rs
