microsandbox_protocol/control/
message.rs1use 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
12pub 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#[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 #[strum(serialize = "control.hello")]
55 Hello,
56 #[strum(serialize = "control.welcome")]
58 Welcome,
59 #[strum(serialize = "control.capabilities")]
61 Capabilities,
62 #[strum(serialize = "control.capabilities.result")]
64 CapabilitiesResult,
65 #[strum(serialize = "control.memory.state")]
67 MemoryState,
68 #[strum(serialize = "control.memory.target")]
70 MemoryTarget,
71 #[strum(serialize = "control.cpu.state")]
73 CpuState,
74 #[strum(serialize = "control.cpu.target")]
76 CpuTarget,
77 #[strum(serialize = "control.secrets.update")]
79 SecretsUpdate,
80 #[strum(serialize = "control.secrets.result")]
82 SecretsResult,
83 #[strum(serialize = "control.error")]
85 Error,
86}
87
88#[derive(Debug, Clone)]
93pub enum ControlOperation {
94 GenerationOne(ControlRequest),
96 CheckpointCreate(CheckpointCreate),
98 DiskCheckpointCreate(DiskCheckpointCreate),
100 BranchCreate(BranchCreate),
102 Pause(Pause),
104 Resume,
106 PauseState,
108 RootDiskGrow(RootDiskGrow),
110 DiskCompact(DiskCompact),
112}
113
114impl ControlMessageType {
119 pub fn as_str(self) -> &'static str {
121 self.into()
122 }
123
124 pub fn from_wire_str(name: &str) -> Option<Self> {
126 name.parse().ok()
127 }
128
129 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 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 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
189pub 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 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 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 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 #[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 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
332impl 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}