protocol: add capture requests and restore activation messages

nvx-rs/src/host/launch.rs

@@ -55,17 +55,7 @@ pub const DEFAULT_HOSTNAME: &str = "nvx-sandbox";
 /// kernel parameters on every command line.
 const HOSTNAME_PARAMETER: &str = "hostname";
 
-/// The identity that workloads run as: a non-root user and group ID.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub struct WorkloadIdentity {
-    /// User ID.
-    pub uid: u32,
-    /// Primary group ID.
-    pub gid: u32,
-    /// Whether the guest creates an account for the IDs if its image lacks a usable one,
-    /// replacing an image account for the UID whose primary group or home does not fit.
-    pub create_account: bool,
-}
+pub use crate::protocol::workload::WorkloadIdentity;
 
 impl WorkloadIdentity {
     /// Alpine's `nobody`, which every guest image has.

nvx-rs/src/host/session.rs

@@ -10,8 +10,8 @@ use super::nvxs::{HostConnection, HostConnectionError};
 use super::platform::{self, Transport};
 use crate::protocol::nvxs::{self, RecordError, RecordType};
 use crate::protocol::{
-    self, APP_CANCEL, APP_ERROR, APP_EXEC, APP_EXIT, APP_FEATURES, APP_HEADER_LEN, APP_MAPS,
-    APP_PING, APP_READY, APP_STDERR, APP_STDOUT, APP_STOP, APP_STOPPED, CAPABILITY_LEN,
+    self, APP_CANCEL, APP_CAPTURE, APP_ERROR, APP_EXEC, APP_EXIT, APP_FEATURES, APP_HEADER_LEN,
+    APP_MAPS, APP_PING, APP_READY, APP_STDERR, APP_STDOUT, APP_STOP, APP_STOPPED, CAPABILITY_LEN,
     ExitCategory, GuestFeatures, HostMap, OUTER_MAX_PAYLOAD, ProtocolError, UNSUPPORTED_OPERATION,
     Workload,
 };
@@ -156,7 +156,29 @@ impl ControlSession {
 
     /// Waits until the guest agent acknowledges readiness.
     pub fn ping(&mut self, deadline: Instant) -> Result<(), SessionError> {
-        self.acknowledged(APP_PING, APP_READY, "readiness", deadline)
+        self.acknowledged(APP_PING, &[], APP_READY, "readiness", deadline)
+    }
+
+    /// Asks nvx-init to capture the VM, killing its executions first when requested.
+    #[cfg_attr(
+        not(test),
+        expect(
+            dead_code,
+            reason = "SL3.7b requests capture once nvx-init handles it in SL7.7a"
+        )
+    )]
+    pub(crate) fn capture(
+        &mut self,
+        kill_executions: bool,
+        deadline: Instant,
+    ) -> Result<(), SessionError> {
+        self.acknowledged(
+            APP_CAPTURE,
+            &protocol::encode_capture_payload(kill_executions),
+            APP_READY,
+            "capture",
+            deadline,
+        )
     }
 
     /// Asks the guest agent which control features its image provides.
@@ -259,7 +281,7 @@ impl ControlSession {
 
     /// Asks the guest agent to shut the VM down.
     pub fn stop(&mut self, deadline: Instant) -> Result<(), SessionError> {
-        self.acknowledged(APP_STOP, APP_STOPPED, "stop", deadline)
+        self.acknowledged(APP_STOP, &[], APP_STOPPED, "stop", deadline)
     }
 
     /// Sends an application request of `kind` with `request_id` and `payload` in the next
@@ -338,17 +360,18 @@ impl ControlSession {
         Ok((kind, status, payload))
     }
 
-    /// Sends a request of `kind` without a payload and waits until the guest agent acknowledges
+    /// Sends a request of `kind` and waits until the guest agent acknowledges
     /// `what` with an empty `reply` of status 0.
     fn acknowledged(
         &mut self,
         kind: u8,
+        payload: &[u8],
         reply: u8,
         what: &str,
         deadline: Instant,
     ) -> Result<(), SessionError> {
         let request_id = request_id()?;
-        self.send_app(kind, request_id, &[], Some(deadline))?;
+        self.send_app(kind, request_id, payload, Some(deadline))?;
         let (kind, response_id, status, payload) = self.read_app(Some(deadline))?;
         if kind != reply || response_id != request_id || status != 0 || !payload.is_empty() {
             return Err(SessionError::Protocol(format!(

nvx-rs/src/host/session/tests.rs

@@ -575,6 +575,40 @@ fn pings_and_stops_need_an_empty_acknowledgement() {
     }
 }
 
+#[test]
+fn capture_sends_both_kill_flags_and_waits_for_acknowledgement() {
+    let (mut session, requests) = attached_recording(|kind, id| {
+        assert_eq!(kind, APP_CAPTURE);
+        assert_ne!(id, 0);
+        encode_app(APP_READY, id, 0, &[])
+    });
+    session.capture(false, deadline()).unwrap();
+    session.capture(true, deadline()).unwrap();
+    assert_eq!(
+        *requests.lock().unwrap(),
+        [(APP_CAPTURE, vec![0]), (APP_CAPTURE, vec![1])]
+    );
+}
+
+#[test]
+fn capture_refuses_invalid_acknowledgements() {
+    let answers: [fn(u8, u64) -> Vec<u8>; 5] = [
+        |_, id| encode_app(APP_ERROR, id, 16, b"busy"),
+        |_, id| encode_app(APP_STOPPED, id, 0, &[]),
+        |_, id| encode_app(APP_READY, id ^ 1, 0, &[]),
+        |_, id| encode_app(APP_READY, id, 1, &[]),
+        |_, id| encode_app(APP_READY, id, 0, b"x"),
+    ];
+    for answer in answers {
+        let error = attached(answer).capture(false, deadline()).unwrap_err();
+        assert_violation(
+            &error,
+            "the guest agent did not acknowledge capture",
+            "an invalid capture acknowledgement",
+        );
+    }
+}
+
 #[test]
 fn features_report_what_the_guest_provides() {
     let mut session = attached(|kind, id| {

nvx-rs/src/protocol/app.rs

@@ -6,6 +6,28 @@
 //! of their request.
 
 use super::{ProtocolError, array, violation};
+use crate::activation::ActivationRecord;
+
+use self::workload::WorkloadIdentity;
+
+/// Shared identity data; host-side selection remains in the host module.
+pub(crate) mod workload {
+    /// The identity that workloads run as: a non-root user and group ID.
+    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
+    #[cfg_attr(
+        not(feature = "host"),
+        expect(unreachable_pub, reason = "public host re-export when host is enabled")
+    )]
+    pub struct WorkloadIdentity {
+        /// User ID.
+        pub uid: u32,
+        /// Primary group ID.
+        pub gid: u32,
+        /// Whether the guest creates an account for the IDs if its image lacks a usable one,
+        /// replacing an image account for the UID whose primary group or home does not fit.
+        pub create_account: bool,
+    }
+}
 
 /// Magic of every application frame.
 pub const APP_MAGIC: [u8; 4] = *b"NVXC";
@@ -28,6 +50,21 @@ pub const APP_FEATURES: u8 = 5;
 /// Request that the guest agent mount entries of the host mapping table; the payload is a `MAPS`
 /// payload.
 pub const APP_MAPS: u8 = 6;
+/// Request that nvx-init capture the VM after optionally killing its executions.
+#[cfg_attr(
+    not(any(test, feature = "host")),
+    expect(dead_code, reason = "SL7.7a makes nvx-init handle capture requests")
+)]
+pub(crate) const APP_CAPTURE: u8 = 7;
+/// Request that nvx-init activate a restored VM with its record and workload identity.
+#[cfg_attr(
+    not(test),
+    expect(
+        dead_code,
+        reason = "SL4.3 sends activation requests to nvx-init in SL7.7b"
+    )
+)]
+pub(crate) const APP_ACTIVATE: u8 = 8;
 /// Response that acknowledges a request.
 pub const APP_READY: u8 = 0x81;
 /// Response that carries a chunk of a workload's standard output.
@@ -112,6 +149,92 @@ pub fn decode_app(frame: &[u8]) -> Result<AppFrame<'_>, ProtocolError> {
     })
 }
 
+/// Encodes the capture request's flag as one byte, 0 or 1.
+#[cfg_attr(
+    not(any(test, feature = "host")),
+    expect(dead_code, reason = "SL7.7a makes nvx-init handle capture requests")
+)]
+pub(crate) fn encode_capture_payload(kill_executions: bool) -> [u8; 1] {
+    [u8::from(kill_executions)]
+}
+
+/// Decodes the capture request's sole kill-executions flag.
+#[cfg_attr(
+    not(test),
+    expect(dead_code, reason = "SL7.7a makes nvx-init handle capture requests")
+)]
+pub(crate) fn decode_capture_payload(payload: &[u8]) -> Result<bool, ProtocolError> {
+    match payload {
+        [0] => Ok(false),
+        [1] => Ok(true),
+        _ => Err(violation(
+            "the capture request must contain one kill-executions flag",
+        )),
+    }
+}
+
+/// Encodes the persistent workload identity, then the existing activation record.
+///
+/// The prefix is the UID and GID as little-endian `u32`s, then the account-creation flag as
+/// exactly 0 or 1. The remainder is the activation record's unchanged encoding.
+#[cfg_attr(
+    not(test),
+    expect(dead_code, reason = "SL4.3 encodes activation requests")
+)]
+pub(crate) fn encode_activate_payload(
+    record: &ActivationRecord,
+    workload: &WorkloadIdentity,
+) -> Result<Vec<u8>, ProtocolError> {
+    check_workload_identity(workload)?;
+    let record = record.encode();
+    let mut payload = Vec::with_capacity(9 + record.len());
+    payload.extend_from_slice(&workload.uid.to_le_bytes());
+    payload.extend_from_slice(&workload.gid.to_le_bytes());
+    payload.push(u8::from(workload.create_account));
+    payload.extend_from_slice(&record);
+    Ok(payload)
+}
+
+/// Decodes the persistent workload identity and the activation record that follows it.
+#[cfg_attr(
+    not(test),
+    expect(dead_code, reason = "SL7.7b makes nvx-init decode activation requests")
+)]
+pub(crate) fn decode_activate_payload(
+    payload: &[u8],
+) -> Result<(ActivationRecord, WorkloadIdentity), ProtocolError> {
+    let prefix = payload
+        .get(..9)
+        .ok_or_else(|| violation("the activate request is truncated"))?;
+    let workload = WorkloadIdentity {
+        uid: u32::from_le_bytes(array(&prefix[..4])),
+        gid: u32::from_le_bytes(array(&prefix[4..8])),
+        create_account: match prefix[8] {
+            0 => false,
+            1 => true,
+            _ => {
+                return Err(violation(
+                    "the activate request has an invalid account flag",
+                ));
+            }
+        },
+    };
+    check_workload_identity(&workload)?;
+    let record = ActivationRecord::decode(&payload[9..])
+        .map_err(|error| violation(format!("invalid activation record: {error}")))?;
+    Ok((record, workload))
+}
+
+/// Matches the guest's non-root credential range, excluding the unchanged-ID sentinel.
+fn check_workload_identity(workload: &WorkloadIdentity) -> Result<(), ProtocolError> {
+    if matches!(workload.uid, 0 | u32::MAX) || matches!(workload.gid, 0 | u32::MAX) {
+        return Err(violation(
+            "activation workload requires a non-root UID and GID below u32::MAX",
+        ));
+    }
+    Ok(())
+}
+
 /// Exit category reported by the guest agent, the payload of an [`APP_EXIT`] response.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum ExitCategory {

nvx-rs/src/protocol/mod.rs

@@ -60,6 +60,8 @@ pub use app::{
     BUSY, CWD_FAILED, ExitCategory, INVALID_REQUEST, LAUNCH_FAILED, MAPPING_FAILED,
     MAPPINGS_INCOMPLETE, UNSUPPORTED_OPERATION, decode_app, encode_app,
 };
+#[cfg(feature = "host")]
+pub(crate) use app::{APP_CAPTURE, encode_capture_payload, workload};
 pub use boot::{MAPS_TOKEN, WORKLOAD_ACCOUNT_TOKEN, parse_host_mapping_count};
 pub use exec::{
     EXEC_CWD_PRESENT, EXEC_ENVIRONMENT_PRESENT, EXEC_EXTENDED, EXEC_INHERIT_DEFAULT_ENV,