microsandbox_control_client/
request.rs1use microsandbox_protocol::{
4 control::{
5 CONTROL_GENERATION, Capabilities, ControlError, CpuState, CpuTarget, Empty, MemoryState,
6 MemoryTarget, SecretChange, SecretsResult,
7 },
8 wire,
9};
10use microsandbox_protocol_client::{EncodedMessage, Message, Request};
11use microsandbox_utils::size::Mebibytes;
12use serde::{Serialize, de::DeserializeOwned};
13
14use crate::{ControlClientError, ControlClientResult, ControlProtocol};
15
16#[derive(Debug, Clone, Copy, Default)]
22pub struct GetCapabilities;
23#[derive(Debug, Clone, Copy, Default)]
25pub struct GetMemoryState;
26#[derive(Debug, Clone, Copy)]
28pub struct SetMemoryTarget {
29 pub total_mib: u64,
31}
32#[derive(Debug, Clone, Copy, Default)]
34pub struct GetCpuState;
35#[derive(Debug, Clone, Copy)]
37pub struct SetCpuTarget {
38 pub online: u32,
40}
41#[derive(Debug, Clone)]
43pub struct UpdateSecrets {
44 pub changes: Vec<SecretChange>,
46}
47
48impl SetMemoryTarget {
53 pub fn new(size: impl Into<Mebibytes>) -> Self {
59 Self {
60 total_mib: u64::from(size.into().as_u32()),
61 }
62 }
63}
64
65impl SetCpuTarget {
66 pub fn new(online: u32) -> Self {
68 Self { online }
69 }
70}
71
72impl UpdateSecrets {
73 pub fn new(changes: Vec<SecretChange>) -> Self {
75 Self { changes }
76 }
77}
78
79impl Request<ControlProtocol> for GetCapabilities {
84 type Response = Capabilities;
85 type Error = ControlClientError;
86 fn message(&self) -> ControlClientResult<EncodedMessage> {
87 prepared("control.capabilities", &Empty {})
88 }
89 fn decode(&self, response: Message) -> ControlClientResult<Self::Response> {
90 checked(response, "control.capabilities.result")
91 }
92}
93
94impl Request<ControlProtocol> for GetMemoryState {
95 type Response = MemoryState;
96 type Error = ControlClientError;
97 fn message(&self) -> ControlClientResult<EncodedMessage> {
98 prepared("control.memory.state", &Empty {})
99 }
100 fn decode(&self, response: Message) -> ControlClientResult<Self::Response> {
101 checked(response, "control.memory.state")
102 }
103}
104
105impl Request<ControlProtocol> for SetMemoryTarget {
106 type Response = MemoryState;
107 type Error = ControlClientError;
108 fn message(&self) -> ControlClientResult<EncodedMessage> {
109 prepared(
110 "control.memory.target",
111 &MemoryTarget {
112 total_mib: self.total_mib,
113 },
114 )
115 }
116 fn decode(&self, response: Message) -> ControlClientResult<Self::Response> {
117 checked(response, "control.memory.state")
118 }
119}
120
121impl Request<ControlProtocol> for GetCpuState {
122 type Response = CpuState;
123 type Error = ControlClientError;
124 fn message(&self) -> ControlClientResult<EncodedMessage> {
125 prepared("control.cpu.state", &Empty {})
126 }
127 fn decode(&self, response: Message) -> ControlClientResult<Self::Response> {
128 checked(response, "control.cpu.state")
129 }
130}
131
132impl Request<ControlProtocol> for SetCpuTarget {
133 type Response = CpuState;
134 type Error = ControlClientError;
135 fn message(&self) -> ControlClientResult<EncodedMessage> {
136 prepared(
137 "control.cpu.target",
138 &CpuTarget {
139 online: self.online,
140 },
141 )
142 }
143 fn decode(&self, response: Message) -> ControlClientResult<Self::Response> {
144 checked(response, "control.cpu.state")
145 }
146}
147
148impl Request<ControlProtocol> for UpdateSecrets {
149 type Response = SecretsResult;
150 type Error = ControlClientError;
151 fn message(&self) -> ControlClientResult<EncodedMessage> {
152 #[derive(Serialize)]
153 struct Payload<'a> {
154 changes: &'a [SecretChange],
155 }
156 prepared(
157 "control.secrets.update",
158 &Payload {
159 changes: &self.changes,
160 },
161 )
162 }
163 fn decode(&self, response: Message) -> ControlClientResult<Self::Response> {
164 checked_with(response, "control.secrets.result", |bytes| {
165 let result = SecretsResult::decode(bytes)?;
166 let valid = match &result {
169 SecretsResult::Complete { applied_count } => {
170 *applied_count as usize == self.changes.len()
171 }
172 SecretsResult::Failed { failed_index, .. } => {
173 (*failed_index as usize) < self.changes.len()
174 }
175 };
176 if !valid {
177 return Err(wire::WireError::InvalidRecord);
178 }
179 Ok(result)
180 })
181 }
182}
183
184fn prepared(name: &str, payload: &impl Serialize) -> ControlClientResult<EncodedMessage> {
189 Ok(EncodedMessage::new(name, wire::encode(payload)?))
190}
191
192fn checked<T: DeserializeOwned>(response: Message, expected: &str) -> ControlClientResult<T> {
193 checked_with(response, expected, wire::decode_record)
194}
195
196fn checked_with<T>(
197 response: Message,
198 expected: &str,
199 decode: impl FnOnce(&[u8]) -> Result<T, wire::WireError>,
200) -> ControlClientResult<T> {
201 if response.v != CONTROL_GENERATION || response.id == 0 || response.flags != 1 {
202 return Err(ControlClientError::InvalidResponse {
203 response: Box::new(response),
204 });
205 }
206 if response.t == "control.error" {
207 let Ok(error) = wire::decode_record::<ControlError>(&response.p) else {
208 return Err(ControlClientError::InvalidResponse {
209 response: Box::new(response),
210 });
211 };
212 return Err(ControlClientError::Peer {
213 error,
214 response: Box::new(response),
215 });
216 }
217 if response.t != expected {
218 return Err(ControlClientError::InvalidResponse {
219 response: Box::new(response),
220 });
221 }
222 decode(&response.p).map_err(|_| ControlClientError::InvalidResponse {
223 response: Box::new(response),
224 })
225}