Skip to main content

microsandbox_protocol/control/
message.rs

1//! Control message inventory and checked dispatch decoding.
2
3use ciborium::Value;
4use serde::{Deserialize, Serialize};
5
6use super::{
7    BranchCreate, CheckpointCreate, ControlRequest, CpuTarget, DiskCheckpointCreate, DiskCompact,
8    Empty, MemoryTarget, Pause, RootDiskGrow, SecretsResult, SecretsUpdate,
9};
10use crate::wire::{Envelope, WireError, decode_value, validate_record};
11
12//--------------------------------------------------------------------------------------------------
13// Constants
14//--------------------------------------------------------------------------------------------------
15
16/// Application messages introduced by framed control generation two.
17///
18/// The inventory is public so other language implementations can pin the same additive surface
19/// without expanding the released, exhaustively matchable [`ControlMessageType`] enum.
20pub const CONTROL_GENERATION_TWO_MESSAGES: &[&str] = &[
21    "control.checkpoint.create",
22    "control.checkpoint.result",
23    "control.disk.checkpoint.create",
24    "control.disk.checkpoint.result",
25    "control.branch.create",
26    "control.branch.result",
27    "control.pause",
28    "control.resume",
29    "control.pause.state",
30    "control.root-disk.grow",
31    "control.root-disk.state",
32    "control.disk.compact",
33    "control.disk.compact.result",
34];
35
36//--------------------------------------------------------------------------------------------------
37// Types
38//--------------------------------------------------------------------------------------------------
39
40/// Known control wire names. The strum spelling is the authoritative mapping.
41#[derive(
42    Debug,
43    Clone,
44    Copy,
45    PartialEq,
46    Eq,
47    Hash,
48    strum::IntoStaticStr,
49    strum::EnumString,
50    strum::EnumIter,
51)]
52pub enum ControlMessageType {
53    /// Initial generation offer.
54    #[strum(serialize = "control.hello")]
55    Hello,
56    /// Selected generation and limits.
57    #[strum(serialize = "control.welcome")]
58    Welcome,
59    /// Facility query.
60    #[strum(serialize = "control.capabilities")]
61    Capabilities,
62    /// Facility response.
63    #[strum(serialize = "control.capabilities.result")]
64    CapabilitiesResult,
65    /// Memory query or terminal memory observation.
66    #[strum(serialize = "control.memory.state")]
67    MemoryState,
68    /// Set a memory target.
69    #[strum(serialize = "control.memory.target")]
70    MemoryTarget,
71    /// CPU query or terminal CPU observation.
72    #[strum(serialize = "control.cpu.state")]
73    CpuState,
74    /// Set a CPU target.
75    #[strum(serialize = "control.cpu.target")]
76    CpuTarget,
77    /// Ordered host secret modifications.
78    #[strum(serialize = "control.secrets.update")]
79    SecretsUpdate,
80    /// Complete or partial secret progress.
81    #[strum(serialize = "control.secrets.result")]
82    SecretsResult,
83    /// Recoverable operation or handshake error.
84    #[strum(serialize = "control.error")]
85    Error,
86}
87
88/// One decoded application operation across all negotiated control generations.
89///
90/// Generation-one's public [`ControlRequest`] remains unchanged so downstream exhaustive matches
91/// keep compiling. Generation-two operations live in this additive dispatch type instead.
92#[derive(Debug, Clone)]
93pub enum ControlOperation {
94    /// A released generation-one operation.
95    GenerationOne(ControlRequest),
96    /// Create one full checkpoint.
97    CheckpointCreate(CheckpointCreate),
98    /// Create one disk-only checkpoint.
99    DiskCheckpointCreate(DiskCheckpointCreate),
100    /// Create one direct local branch without descriptor transfer.
101    BranchCreate(BranchCreate),
102    /// Pause, optionally with an explicit guest-writeback policy.
103    Pause(Pause),
104    /// Resume a resident pause.
105    Resume,
106    /// Inspect resident pause state.
107    PauseState,
108    /// Grow the root disk and filesystem.
109    RootDiskGrow(RootDiskGrow),
110    /// Compact selected owned disk chains.
111    DiskCompact(DiskCompact),
112}
113
114//--------------------------------------------------------------------------------------------------
115// Methods
116//--------------------------------------------------------------------------------------------------
117
118impl ControlMessageType {
119    /// Stable wire spelling.
120    pub fn as_str(self) -> &'static str {
121        self.into()
122    }
123
124    /// Resolve a known spelling, retaining unknown names in the envelope layer.
125    pub fn from_wire_str(name: &str) -> Option<Self> {
126        name.parse().ok()
127    }
128
129    /// Whether this is an application request permitted after the handshake.
130    pub fn is_request(self) -> bool {
131        matches!(
132            self,
133            Self::Capabilities
134                | Self::MemoryState
135                | Self::MemoryTarget
136                | Self::CpuState
137                | Self::CpuTarget
138                | Self::SecretsUpdate
139        )
140    }
141}
142
143impl ControlOperation {
144    /// Decode one request admitted by the negotiated generation.
145    pub fn from_envelope(envelope: &Envelope, generation: u8) -> Result<Self, WireError> {
146        if envelope.v != generation {
147            return Err(WireError::InvalidRecord);
148        }
149        let operation = match envelope.t.as_str() {
150            "control.checkpoint.create" if generation >= 2 => {
151                Self::CheckpointCreate(envelope.payload()?)
152            }
153            "control.disk.checkpoint.create" if generation >= 2 => {
154                Self::DiskCheckpointCreate(envelope.payload()?)
155            }
156            "control.branch.create" if generation >= 2 => Self::BranchCreate(envelope.payload()?),
157            "control.pause" if generation >= 2 => Self::Pause(envelope.payload()?),
158            "control.resume" if generation >= 2 => {
159                envelope.payload::<Empty>()?;
160                Self::Resume
161            }
162            "control.pause.state" if generation >= 2 => {
163                envelope.payload::<Empty>()?;
164                Self::PauseState
165            }
166            "control.root-disk.grow" if generation >= 2 => Self::RootDiskGrow(envelope.payload()?),
167            "control.disk.compact" if generation >= 2 => {
168                // The selector is itself a tagged record. Validate it before serde can collapse
169                // duplicate keys and choose a different mutation target than the wire expressed.
170                let value = decode_value(&envelope.p)?;
171                validate_record(&value)?;
172                let Value::Map(fields) = &value else {
173                    unreachable!()
174                };
175                if let Some((_, target)) = fields
176                    .iter()
177                    .find(|(key, _)| key.as_text() == Some("target"))
178                {
179                    validate_record(target)?;
180                }
181                Self::DiskCompact(value.deserialized().map_err(|_| WireError::InvalidRecord)?)
182            }
183            _ => Self::GenerationOne(ControlRequest::from_envelope(envelope)?),
184        };
185        Ok(operation)
186    }
187}
188
189/// Generation in which a known application message first became available.
190///
191/// Unknown extension names remain caller-owned and therefore return `None`.
192pub fn control_message_min_generation(name: &str) -> Option<u8> {
193    if matches!(name, "control.jobs" | "control.jobs.result") {
194        return Some(2);
195    }
196    if CONTROL_GENERATION_TWO_MESSAGES.contains(&name) {
197        return Some(2);
198    }
199    Some(match name {
200        "control.capabilities"
201        | "control.capabilities.result"
202        | "control.memory.state"
203        | "control.memory.target"
204        | "control.cpu.state"
205        | "control.cpu.target"
206        | "control.secrets.update"
207        | "control.secrets.result"
208        | "control.error" => 1,
209        _ => return None,
210    })
211}
212
213impl ControlRequest {
214    /// Decode a known application request after the server validates its frame.
215    pub fn from_envelope(envelope: &Envelope) -> Result<Self, WireError> {
216        Ok(match ControlMessageType::from_wire_str(&envelope.t) {
217            Some(ControlMessageType::Capabilities) => {
218                envelope.payload::<Empty>()?;
219                Self::Capabilities
220            }
221            Some(ControlMessageType::MemoryState) => {
222                envelope.payload::<Empty>()?;
223                Self::MemoryState
224            }
225            Some(ControlMessageType::CpuState) => {
226                envelope.payload::<Empty>()?;
227                Self::CpuState
228            }
229            Some(ControlMessageType::MemoryTarget) => {
230                let payload: MemoryTarget = envelope.payload()?;
231                Self::MemoryTarget {
232                    total_mib: payload.total_mib,
233                }
234            }
235            Some(ControlMessageType::CpuTarget) => {
236                let payload: CpuTarget = envelope.payload()?;
237                Self::CpuTarget {
238                    online: payload.online,
239                }
240            }
241            Some(ControlMessageType::SecretsUpdate) => {
242                // Serde's tagged-enum buffering alone does not reject every
243                // duplicate. Check each entry before any host operation runs.
244                let value = decode_value(&envelope.p)?;
245                validate_record(&value)?;
246                let Value::Map(fields) = &value else {
247                    unreachable!()
248                };
249                if let Some((_, Value::Array(changes))) = fields
250                    .iter()
251                    .find(|(key, _)| key.as_text() == Some("changes"))
252                {
253                    for change in changes {
254                        validate_record(change)?;
255                    }
256                }
257                let payload: SecretsUpdate =
258                    value.deserialized().map_err(|_| WireError::InvalidRecord)?;
259                Self::SecretsUpdate {
260                    changes: payload.changes,
261                }
262            }
263            _ => return Err(WireError::InvalidRecord),
264        })
265    }
266
267    /// Encode the CBOR operation corresponding to this legacy-compatible value.
268    pub fn envelope(&self, generation: u8) -> Result<Envelope, WireError> {
269        match self {
270            Self::Capabilities => Envelope::new(
271                generation,
272                ControlMessageType::Capabilities.as_str(),
273                &Empty {},
274            ),
275            Self::MemoryState => Envelope::new(
276                generation,
277                ControlMessageType::MemoryState.as_str(),
278                &Empty {},
279            ),
280            Self::CpuState => {
281                Envelope::new(generation, ControlMessageType::CpuState.as_str(), &Empty {})
282            }
283            Self::MemoryTarget { total_mib } => Envelope::new(
284                generation,
285                ControlMessageType::MemoryTarget.as_str(),
286                &MemoryTarget {
287                    total_mib: *total_mib,
288                },
289            ),
290            Self::CpuTarget { online } => Envelope::new(
291                generation,
292                ControlMessageType::CpuTarget.as_str(),
293                &CpuTarget { online: *online },
294            ),
295            Self::SecretsUpdate { changes } => {
296                // Borrow secret entries instead of cloning their plaintext.
297                #[derive(Serialize)]
298                struct Payload<'a> {
299                    changes: &'a [super::SecretChange],
300                }
301                Envelope::new(
302                    generation,
303                    ControlMessageType::SecretsUpdate.as_str(),
304                    &Payload { changes },
305                )
306            }
307        }
308    }
309}
310
311impl SecretsResult {
312    /// Decode sequential progress, checking nested error keys and index equality.
313    pub fn decode(payload: &[u8]) -> Result<Self, WireError> {
314        let value = decode_value(payload)?;
315        validate_record(&value)?;
316        if let Value::Map(fields) = &value
317            && let Some((_, error)) = fields
318                .iter()
319                .find(|(key, _)| key.as_text() == Some("error"))
320        {
321            validate_record(error)?;
322        }
323        let result: Self = value.deserialized().map_err(|_| WireError::InvalidRecord)?;
324        if matches!(&result, Self::Failed { applied_count, failed_index, .. } if applied_count != failed_index)
325        {
326            return Err(WireError::InvalidRecord);
327        }
328        Ok(result)
329    }
330}
331
332//--------------------------------------------------------------------------------------------------
333// Trait Implementations
334//--------------------------------------------------------------------------------------------------
335
336impl AsRef<str> for ControlMessageType {
337    fn as_ref(&self) -> &str {
338        self.as_str()
339    }
340}
341
342impl Serialize for ControlMessageType {
343    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
344        serializer.serialize_str(self.as_str())
345    }
346}
347
348impl<'de> Deserialize<'de> for ControlMessageType {
349    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
350        let name = String::deserialize(deserializer)?;
351        Self::from_wire_str(&name)
352            .ok_or_else(|| serde::de::Error::custom("unknown control message type"))
353    }
354}