arcbox_computer_runtime/agent/vm_proto/
client.rs

1//! The [`GuestAgent`] port over the `arcbox-vm-proto` channels.
2//!
3//! Nothing here decides anything: it is the mapping from the port's verbs
4//! onto the protocol functions in the sibling modules, plus the three
5//! readiness shapes [`Readiness`] names, each armed on the driver port's
6//! vsock capabilities.
7
8use std::sync::Arc;
9use std::time::Duration;
10
11use arcbox_vm_driver::net::NetworkIdentity;
12use arcbox_vm_driver::{PreparedVm, VmHandle, Vsock, VsockListener};
13use async_trait::async_trait;
14use tokio::sync::mpsc;
15
16use super::files;
17use super::{HandleVsock, READY_PORT, clock, connect_to_port, exec, net, wait_port, wait_ready};
18use crate::agent::{
19    ClockSync, DirWatch, ExecInput, GuestAgent, GuestAgentFactory, GuestFiles, OutputChunk,
20    PortWait, ProbeGate, Readiness, ReadyGate, StartCommand,
21};
22use crate::boot_proto::NetReconfigCommand;
23use crate::error::Result;
24use crate::file_proto::FileStatDto;
25
26/// The guest agent of one running VM, reached over its vsock device.
27pub struct VmProtoAgent {
28    vsock: Arc<dyn Vsock>,
29}
30
31impl VmProtoAgent {
32    /// The agent inside `handle`'s guest.
33    ///
34    /// The port lends the vsock capability out of the handle
35    /// (`VmHandle::vsock` borrows) while callers hold the agent across
36    /// awaits and hand it to detached tasks, so the agent owns the handle.
37    #[must_use]
38    pub fn over(handle: Arc<dyn VmHandle>) -> Self {
39        Self {
40            vsock: Arc::new(HandleVsock(handle)),
41        }
42    }
43}
44
45#[async_trait]
46impl GuestAgent for VmProtoAgent {
47    async fn run(&self, start: StartCommand) -> Result<mpsc::Receiver<Result<OutputChunk>>> {
48        exec::run(self.vsock.as_ref(), start).await
49    }
50
51    async fn exec(
52        &self,
53        start: StartCommand,
54    ) -> Result<(ExecInput, mpsc::Receiver<Result<OutputChunk>>)> {
55        exec::exec(self.vsock.as_ref(), start).await
56    }
57
58    async fn sync_clock(&self) -> Result<ClockSync> {
59        clock::sync_clock(self.vsock.as_ref()).await
60    }
61
62    async fn reconfigure_network(&self, cmd: &NetReconfigCommand) -> Result<()> {
63        net::reconfigure_network(self.vsock.as_ref(), cmd).await
64    }
65
66    async fn wait_for_port(&self, port: u16, timeout: Duration) -> Result<PortWait> {
67        wait_port::wait_for_port(self.vsock.as_ref(), port, timeout).await
68    }
69
70    fn files(&self) -> &dyn GuestFiles {
71        self
72    }
73}
74
75#[async_trait]
76impl GuestFiles for VmProtoAgent {
77    async fn read(&self, path: &str) -> Result<Vec<u8>> {
78        files::read_file(self.vsock.as_ref(), path).await
79    }
80
81    async fn write(&self, path: &str, mode: u32, data: &[u8]) -> Result<()> {
82        files::write_file(self.vsock.as_ref(), path, mode, data).await
83    }
84
85    async fn stat(&self, path: &str) -> Result<FileStatDto> {
86        files::stat_file(self.vsock.as_ref(), path).await
87    }
88
89    async fn list(&self, path: &str) -> Result<Vec<FileStatDto>> {
90        files::list_dir(self.vsock.as_ref(), path).await
91    }
92
93    async fn make_dir(&self, path: &str, mode: u32) -> Result<()> {
94        files::make_dir(self.vsock.as_ref(), path, mode).await
95    }
96
97    async fn remove(&self, path: &str, recursive: bool) -> Result<()> {
98        files::remove_entry(self.vsock.as_ref(), path, recursive).await
99    }
100
101    async fn rename(&self, from: &str, to: &str) -> Result<()> {
102        files::move_entry(self.vsock.as_ref(), from, to).await
103    }
104
105    async fn watch(&self, path: &str, recursive: bool) -> Result<DirWatch> {
106        files::watch_dir(self.vsock.as_ref(), path, recursive).await
107    }
108}
109
110/// Builds [`VmProtoAgent`]s, and knows how their guests announce readiness.
111///
112/// The default is what `vm-agent` does: dial the host back on
113/// [`READY_PORT`] once its exec and file listeners are up. The other two
114/// [`Readiness`] shapes are here because the port promises them to
115/// implementations whose guest cannot dial out — a composer selects one
116/// with [`Self::with_readiness`].
117pub struct VmProtoAgentFactory {
118    readiness: Readiness,
119}
120
121impl Default for VmProtoAgentFactory {
122    fn default() -> Self {
123        Self {
124            readiness: Readiness::DialOut { port: READY_PORT },
125        }
126    }
127}
128
129impl VmProtoAgentFactory {
130    /// A factory whose guests announce themselves the given way.
131    #[must_use]
132    pub const fn with_readiness(readiness: Readiness) -> Self {
133        Self { readiness }
134    }
135}
136
137#[async_trait]
138impl GuestAgentFactory for VmProtoAgentFactory {
139    fn readiness(&self) -> Readiness {
140        self.readiness
141    }
142
143    async fn arm_readiness(&self, prepared: &dyn PreparedVm) -> Result<Box<dyn ReadyGate>> {
144        match self.readiness {
145            // Bind the readiness listener BEFORE the guest starts: vm-agent
146            // dials host port READY_PORT as soon as it is serving, and the
147            // VMM forwards that guest-initiated connect only if someone is
148            // already listening — otherwise the guest is reset and the one
149            // readiness event is lost. The prepared VM binds it where the
150            // guest's dial-out lands.
151            Readiness::DialOut { port } => {
152                let listener = match prepared.vsock_listener() {
153                    Some(listen) => listen.listen(port).await,
154                    None => Err(arcbox_vm_driver::Error::InvalidSpec(
155                        "the vm driver cannot listen for the guest's readiness dial-out".into(),
156                    )),
157                }?;
158                Ok(Box::new(DialOutGate { listener }))
159            }
160            Readiness::Poll { port } => Ok(Box::new(PollGate { port })),
161            // Transport-independent: it round-trips the agent, which every
162            // implementation has.
163            Readiness::Probe => Ok(Box::new(ProbeGate)),
164        }
165    }
166
167    fn connect(
168        &self,
169        handle: Arc<dyn VmHandle>,
170        _net: Option<&NetworkIdentity>,
171    ) -> Result<Arc<dyn GuestAgent>> {
172        Ok(Arc::new(VmProtoAgent::over(handle)))
173    }
174}
175
176/// [`Readiness::DialOut`]: accept the guest's one connection on the
177/// pre-bound listener and read the single byte it writes.
178struct DialOutGate {
179    listener: Box<dyn VsockListener>,
180}
181
182#[async_trait]
183impl ReadyGate for DialOutGate {
184    async fn wait(
185        &mut self,
186        _handle: &Arc<dyn VmHandle>,
187        _agent: &Arc<dyn GuestAgent>,
188    ) -> Result<()> {
189        wait_ready(&mut *self.listener).await
190    }
191}
192
193/// [`Readiness::Poll`]: dial the guest until its agent accepts.
194///
195/// `connect_to_port`'s retry set is what makes this a wait rather than a
196/// race: a guest that is not listening yet answers `ConnectionRefused`,
197/// the one transient outcome, and a VM that never gets there ends the wait
198/// on that budget.
199struct PollGate {
200    port: u32,
201}
202
203#[async_trait]
204impl ReadyGate for PollGate {
205    async fn wait(
206        &mut self,
207        handle: &Arc<dyn VmHandle>,
208        _agent: &Arc<dyn GuestAgent>,
209    ) -> Result<()> {
210        let vsock = HandleVsock(Arc::clone(handle));
211        connect_to_port(&vsock, self.port).await.map(drop)
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use std::io::Write as _;
218
219    use arcbox_vm_driver::testkit::FakeDriver;
220    use arcbox_vm_driver::{
221        BootSpec, ConsoleSpec, IsolationSpec, VmDriver as _, VmId, VmSpec, VsockSpec,
222    };
223
224    use super::*;
225
226    fn vm_spec(id: &VmId) -> VmSpec {
227        VmSpec {
228            id: id.clone(),
229            cpus: 1,
230            memory_mib: 128,
231            boot: BootSpec::Kernel {
232                image: "/vmlinux".into(),
233                cmdline: String::new(),
234                initrd: None,
235            },
236            disks: vec![],
237            nics: vec![],
238            vsock: Some(VsockSpec { guest_cid: 3 }),
239            shares: vec![],
240            console: ConsoleSpec::Off,
241            balloon: false,
242            entropy: false,
243            dirty_tracking: false,
244            isolation: IsolationSpec::None,
245        }
246    }
247
248    /// Arm a gate on a prepared VM, boot it, and return everything the
249    /// wait needs.
250    async fn armed(
251        readiness: Readiness,
252        dir: &std::path::Path,
253    ) -> (
254        FakeDriver,
255        VmId,
256        Arc<dyn VmHandle>,
257        Arc<dyn GuestAgent>,
258        Box<dyn ReadyGate>,
259    ) {
260        let driver = FakeDriver::new();
261        let id = VmId::new("gate").unwrap();
262        let prepared = driver
263            .prepare()
264            .unwrap()
265            .prepare(&id, &IsolationSpec::None, dir)
266            .await
267            .unwrap();
268        let factory = VmProtoAgentFactory::with_readiness(readiness);
269        let gate = factory.arm_readiness(prepared.as_ref()).await.unwrap();
270        let handle: Arc<dyn VmHandle> = Arc::from(prepared.boot(vm_spec(&id)).await.unwrap());
271        let agent = factory.connect(Arc::clone(&handle), None).unwrap();
272        (driver, id, handle, agent, gate)
273    }
274
275    /// The dial-out gate is the guest's one connection plus its one byte,
276    /// on a listener that was bound before the guest existed.
277    #[tokio::test]
278    async fn dial_out_gate_takes_the_guests_one_byte() {
279        let dir = tempfile::tempdir().unwrap();
280        let (driver, id, handle, agent, mut gate) =
281            armed(Readiness::DialOut { port: READY_PORT }, dir.path()).await;
282
283        let mut guest = driver.guest_dial(&id, READY_PORT).unwrap();
284        guest.write_all(&[1]).unwrap();
285
286        gate.wait(&handle, &agent).await.unwrap();
287    }
288
289    /// The poll gate reaches the guest through the booted VM — it had
290    /// nothing to bind when it was armed.
291    #[tokio::test]
292    async fn poll_gate_dials_the_booted_guest() {
293        let dir = tempfile::tempdir().unwrap();
294        let (_driver, _id, handle, agent, mut gate) =
295            armed(Readiness::Poll { port: 52 }, dir.path()).await;
296
297        gate.wait(&handle, &agent).await.unwrap();
298    }
299}
arcbox_computer_runtime/agent/vm_proto/client.rs

arcbox_computer_runtime/agent/vm_proto/
client.rs
