diff --git a/nvx-rs/src/host/cli/session.rs b/nvx-rs/src/host/cli/session.rs index 3cacb01a8..681757b90 100644 --- a/nvx-rs/src/host/cli/session.rs +++ b/nvx-rs/src/host/cli/session.rs @@ -269,6 +269,33 @@ impl Session { received.map_err(|error| self.fail(error)) } + /// Reads the next response, which must answer `request_id`, and returns its kind, status, + /// and payload. + fn reply(&mut self, request_id: u64, deadline: Deadline) -> Result<(u8, i32, Vec)> { + let (kind, response_id, status, payload) = self.receive(deadline)?; + if response_id != request_id { + return Err(Exit::error( + "managed guest returned a mismatched request ID", + )); + } + Ok((kind, status, payload)) + } + + /// Sends a request of `kind` without a payload and waits up to `timeout` seconds until the + /// guest agent acknowledges `what` with an empty `reply` of status 0. + fn acknowledged(&mut self, kind: u8, reply: u8, what: &str, timeout: f64) -> Result { + let request_id = request_id()?; + let deadline = Deadline::after(timeout); + self.send(kind, request_id, &[], deadline)?; + let (kind, response_id, status, payload) = self.receive(deadline)?; + if kind != reply || response_id != request_id || status != 0 || !payload.is_empty() { + return Err(Exit::error(format!( + "managed guest did not acknowledge {what}" + ))); + } + Ok(()) + } + /// Closes the session after a failed transfer and returns the failure. fn fail(&mut self, error: SessionError) -> Exit { self.inner = None; @@ -308,14 +335,7 @@ impl Session { /// Waits until the guest agent acknowledges readiness, as `ControlSession.ping` does. pub(crate) fn ping(&mut self, timeout: f64) -> Result { - let request_id = request_id()?; - let deadline = Deadline::after(timeout); - self.send(APP_PING, request_id, &[], deadline)?; - let (kind, response_id, status, payload) = self.receive(deadline)?; - if kind != APP_READY || response_id != request_id || status != 0 || !payload.is_empty() { - return Err(Exit::error("managed guest did not acknowledge readiness")); - } - Ok(()) + self.acknowledged(APP_PING, APP_READY, "readiness", timeout) } /// Asks the guest agent which control features its image provides. An agent that predates @@ -324,12 +344,7 @@ impl Session { let request_id = request_id()?; let deadline = Deadline::after(timeout); self.send(APP_FEATURES, request_id, &[], deadline)?; - let (kind, response_id, status, payload) = self.receive(deadline)?; - if response_id != request_id { - return Err(Exit::error( - "managed guest returned a mismatched request ID", - )); - } + let (kind, status, payload) = self.reply(request_id, deadline)?; match kind { APP_READY if status == 0 => GuestFeatures::decode(&payload) .map_err(|_| Exit::error("managed guest returned an invalid features response")), @@ -406,12 +421,7 @@ impl Session { let mut stdout = Vec::new(); let mut stderr = Vec::new(); loop { - let (kind, response_id, status, response) = self.receive(deadline)?; - if response_id != request_id { - return Err(Exit::error( - "managed guest returned a mismatched request ID", - )); - } + let (kind, status, response) = self.reply(request_id, deadline)?; match kind { APP_STDOUT | APP_STDERR => { // The guest agent forwards at most this much output, so more breaks the @@ -474,14 +484,7 @@ impl Session { /// Asks the guest agent to shut the VM down, as `ControlSession.stop` does. pub(crate) fn stop(&mut self, timeout: f64) -> Result { - let request_id = request_id()?; - let deadline = Deadline::after(timeout); - self.send(APP_STOP, request_id, &[], deadline)?; - let (kind, response_id, status, payload) = self.receive(deadline)?; - if kind != APP_STOPPED || response_id != request_id || status != 0 || !payload.is_empty() { - return Err(Exit::error("managed guest did not acknowledge stop")); - } - Ok(()) + self.acknowledged(APP_STOP, APP_STOPPED, "stop", timeout) } } diff --git a/nvx-rs/src/host/cli/session/tests.rs b/nvx-rs/src/host/cli/session/tests.rs index 109d0221c..588ae3ba6 100644 --- a/nvx-rs/src/host/cli/session/tests.rs +++ b/nvx-rs/src/host/cli/session/tests.rs @@ -808,6 +808,15 @@ fn protocol_violations_keep_the_messages_of_python() { "managed guest returned a mismatched request ID" ); worker.join().unwrap(); + let (mut session, worker) = session_with(|guest| { + let (_, request_id, _) = guest.request(); + guest.send(APP_READY, request_id ^ 1, 0, &GuestFeatures::NONE.encode()); + }); + assert_eq!( + message(session.features(5.0)), + "managed guest returned a mismatched request ID" + ); + worker.join().unwrap(); let (mut session, worker) = session_with(|guest| { let (_, request_id, _) = guest.request(); guest.send(APP_READY, request_id, 0, b""); @@ -817,6 +826,18 @@ fn protocol_violations_keep_the_messages_of_python() { "managed guest did not acknowledge stop" ); worker.join().unwrap(); + // An acknowledgement has status 0 and no payload. + for (status, payload) in [(1, &b""[..]), (0, b"x")] { + let (mut session, worker) = session_with(move |guest| { + let (_, request_id, _) = guest.request(); + guest.send(APP_READY, request_id, status, payload); + }); + assert_eq!( + message(session.ping(5.0)), + "managed guest did not acknowledge readiness" + ); + worker.join().unwrap(); + } } /// The endpoints of POSIX hosts: Unix-domain sockets that this process serves.