1mod payload {
13 use serde::{Deserialize, Deserializer, Serializer};
14
15 pub fn serialize<S: Serializer>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error> {
16 serializer.serialize_bytes(bytes)
17 }
18
19 pub fn deserialize<'de, D: Deserializer<'de>>(deserializer: D) -> Result<&'de [u8], D::Error> {
20 <&[u8]>::deserialize(deserializer)
21 }
22}
23
24use serde::{Deserialize, Serialize};
32use super::{WireAnswer, WireChild, WireChildAction, WireChildCall, WireChildControl, WireCommand};
34
35const VERSION: u8 = 6;
38pub(super) const LIMIT: usize = 16 * 1024 * 1024;
39
40#[derive(Serialize, Deserialize)]
47struct Value<'a> {
48 key: &'a str,
49
50 #[serde(borrow, with = "payload")]
51 payload: &'a [u8],
52}
53
54#[derive(Debug, thiserror::Error)]
57pub enum WireCodecError {
58 #[error("operation envelope exceeds the 16 MiB limit")]
59 Oversized,
60
61 #[error("unsupported operation envelope version")]
62 Version,
63
64 #[error("unexpected operation envelope kind")]
65 Kind,
66
67 #[error("trailing operation envelope bytes")]
68 Trailing,
69
70 #[error("invalid operation envelope")]
71 Encoding(#[from] postcard::Error),
72}
73
74#[derive(Serialize, Deserialize)]
77enum Envelope<'a> {
78 Prepare {
79 invocation: u64,
80 operation: &'a str,
81
82 #[serde(borrow, with = "payload")]
83 payload: &'a [u8],
84 },
85 Initialize {
86 invocation: u64,
87 context: [u8; 32],
88 },
89 Attempt {
90 invocation: u64,
91 context: [u8; 32],
92
93 #[serde(borrow)]
94 values: Vec<Value<'a>>,
95 },
96 Finalize {
97 invocation: u64,
98 },
99 Ready {
100 #[serde(borrow, with = "payload")]
101 payload: &'a [u8],
102 },
103 Failed {
104 failure: crate::Failure,
105 },
106 Child {
107 context: [u8; 32],
108 generation_context: Option<[u8; 32]>,
109
110 #[serde(borrow)]
111 values: Vec<Value<'a>>,
112
113 operation: &'a str,
114
115 #[serde(borrow, with = "payload")]
116 payload: &'a [u8],
117
118 submit: bool,
119 },
120 PrepareGrpc {
121 invocation: u64,
122 operation: &'a str,
123
124 #[serde(borrow, with = "payload")]
125 payload: &'a [u8],
126 },
127 StreamAttempt {
128 invocation: u64,
129 epoch: u64,
130 context: [u8; 32],
131
132 #[serde(borrow)]
133 values: Vec<Value<'a>>,
134 },
135 StreamSend {
136 invocation: u64,
137 epoch: u64,
138
139 #[serde(borrow, with = "payload")]
140 payload: &'a [u8],
141 },
142 StreamFinish {
143 invocation: u64,
144 epoch: u64,
145 },
146 StreamReceive {
147 invocation: u64,
148 epoch: u64,
149 },
150 StreamItem {
151 epoch: u64,
152 transferred: bool,
153
154 #[serde(borrow, with = "payload")]
155 payload: &'a [u8],
156 },
157 StreamEnd {
158 epoch: u64,
159 transferred: bool,
160 },
161 ChildControl {
162 context: [u8; 32],
163 instance: u64,
164 action: WireChildAction,
165 },
166}
167
168impl WireCommand {
173 pub fn prepare_encoded_len(invocation: u64, operation: &str, payload: usize) -> Result<usize, WireCodecError> {
175 payload_size(Envelope::Prepare { invocation, operation, payload: &[] }, payload)
176 }
177
178 pub fn encode(&self) -> Result<Vec<u8>, WireCodecError> {
179 encode(self.envelope())
180 }
181
182 pub fn decode(bytes: &[u8]) -> Result<Self, WireCodecError> {
183 Ok(match decode(bytes)? {
184 Envelope::Prepare { invocation, operation, payload } => {
185 Self::Prepare { invocation, operation: operation.to_owned(), payload: payload.to_vec() }
186 }
187 Envelope::PrepareGrpc { invocation, operation, payload } => {
188 Self::PrepareGrpc { invocation, operation: operation.to_owned(), payload: payload.to_vec() }
189 }
190 Envelope::Initialize { invocation, context } => Self::Initialize { invocation, context },
191 Envelope::Attempt { invocation, context, values } => Self::Attempt { invocation, context, values: owned(values)? },
192 Envelope::Finalize { invocation } => Self::Finalize { invocation },
193 Envelope::StreamAttempt { invocation, epoch, context, values } => {
194 Self::StreamAttempt { invocation, epoch, context, values: owned(values)? }
195 }
196 Envelope::StreamSend { invocation, epoch, payload } => Self::StreamSend { invocation, epoch, payload: payload.to_vec() },
197 Envelope::StreamFinish { invocation, epoch } => Self::StreamFinish { invocation, epoch },
198 Envelope::StreamReceive { invocation, epoch } => Self::StreamReceive { invocation, epoch },
199 _ => return Err(WireCodecError::Kind),
200 })
201 }
202
203 fn envelope(&self) -> Envelope<'_> {
204 match self {
205 Self::Prepare { invocation, operation, payload } => Envelope::Prepare { invocation: *invocation, operation, payload },
206 Self::PrepareGrpc { invocation, operation, payload } => Envelope::PrepareGrpc { invocation: *invocation, operation, payload },
207 Self::Initialize { invocation, context } => Envelope::Initialize { invocation: *invocation, context: *context },
208 Self::Attempt { invocation, context, values } => {
209 Envelope::Attempt { invocation: *invocation, context: *context, values: borrowed(values) }
210 }
211 Self::StreamAttempt { invocation, epoch, context, values } => {
212 Envelope::StreamAttempt { invocation: *invocation, epoch: *epoch, context: *context, values: borrowed(values) }
213 }
214 Self::StreamSend { invocation, epoch, payload } => Envelope::StreamSend { invocation: *invocation, epoch: *epoch, payload },
215 Self::StreamFinish { invocation, epoch } => Envelope::StreamFinish { invocation: *invocation, epoch: *epoch },
216 Self::StreamReceive { invocation, epoch } => Envelope::StreamReceive { invocation: *invocation, epoch: *epoch },
217 Self::Finalize { invocation } => Envelope::Finalize { invocation: *invocation },
218 }
219 }
220}
221
222impl WireAnswer {
223 pub fn ready_encoded_len(payload: usize) -> Result<usize, WireCodecError> {
225 payload_size(Envelope::Ready { payload: &[] }, payload)
226 }
227
228 pub fn encode(&self) -> Result<Vec<u8>, WireCodecError> {
229 let envelope = match self {
230 Self::Ready { payload } => Envelope::Ready { payload },
231 Self::Stream { epoch, transferred, payload: Some(payload) } => {
232 Envelope::StreamItem { epoch: *epoch, transferred: *transferred, payload }
233 }
234 Self::Stream { epoch, transferred, payload: None } => Envelope::StreamEnd { epoch: *epoch, transferred: *transferred },
235 Self::Failed { failure } => Envelope::Failed { failure: failure.clone() },
236 };
237 encode(envelope)
238 }
239
240 pub fn decode(bytes: &[u8]) -> Result<Self, WireCodecError> {
241 Ok(match decode(bytes)? {
242 Envelope::StreamItem { epoch, transferred, payload } => Self::Stream { epoch, transferred, payload: Some(payload.to_vec()) },
243 Envelope::StreamEnd { epoch, transferred } => Self::Stream { epoch, transferred, payload: None },
244 Envelope::Ready { payload } => Self::Ready { payload: payload.to_vec() },
245 Envelope::Failed { failure } => Self::Failed { failure },
246 _ => return Err(WireCodecError::Kind),
247 })
248 }
249}
250
251impl WireChildCall {
252 pub fn payload_encoded_len(
255 operation: &str, payload: usize, generation_context: bool, values: &[super::super::CallValue],
256 ) -> Result<usize, WireCodecError> {
257 check_values(values)?;
258 let envelope = Envelope::Child {
259 context: [0; 32],
260 generation_context: generation_context.then_some([0; 32]),
261 values: Vec::new(),
262 operation,
263 payload: &[],
264 submit: false,
265 };
266 let count = postcard::experimental::serialized_size(&values.len())?;
267 let mut size = add_size(payload_size(envelope, payload)?, count - 1)?;
268 for value in values {
269 let value = Value { key: &value.key, payload: &value.payload };
270 size = add_size(size, postcard::experimental::serialized_size(&value)?)?;
271 }
272 Ok(size)
273 }
274
275 pub fn encode(&self) -> Result<Vec<u8>, WireCodecError> {
276 check_values(&self.values)?;
277 encode(Envelope::Child {
278 context: self.context,
279 generation_context: self.generation_context,
280 values: borrowed(&self.values),
281 operation: &self.operation,
282 payload: &self.payload,
283 submit: self.submit,
284 })
285 }
286
287 pub fn decode(bytes: &[u8]) -> Result<Self, WireCodecError> {
288 match WireChild::decode(bytes)? {
289 WireChild::Call(call) => Ok(call),
290 WireChild::Control(_) => Err(WireCodecError::Kind),
291 }
292 }
293}
294
295impl WireChildControl {
296 pub fn encode(&self) -> Result<Vec<u8>, WireCodecError> {
297 encode(Envelope::ChildControl { context: self.context, instance: self.instance, action: self.action })
298 }
299}
300
301impl WireChild {
302 pub fn decode(bytes: &[u8]) -> Result<Self, WireCodecError> {
303 let child = match decode(bytes)? {
304 Envelope::Child { context, generation_context, values, operation, payload, submit } => Self::Call(WireChildCall {
305 context,
306 generation_context,
307 values: owned(values)?,
308 operation: operation.to_owned(),
309 payload: payload.to_vec(),
310 submit,
311 }),
312 Envelope::ChildControl { context, instance, action } => Self::Control(WireChildControl { context, instance, action }),
313 _ => return Err(WireCodecError::Kind),
314 };
315 Ok(child)
316 }
317
318 pub fn context(&self) -> &[u8; 32] {
319 match self {
320 Self::Call(call) => &call.context,
321 Self::Control(control) => &control.context,
322 }
323 }
324}
325
326fn check_values(values: &[super::super::CallValue]) -> Result<(), WireCodecError> {
331 let bytes = values.iter().map(|value| value.key.len().saturating_add(value.payload.len())).fold(0usize, usize::saturating_add);
332 if values.len() > super::super::PORTABLE_VALUES || bytes > super::super::PORTABLE_BYTES {
333 return Err(WireCodecError::Oversized);
334 }
335 for (index, value) in values.iter().enumerate() {
336 if values[..index].iter().any(|previous| previous.key == value.key) {
337 return Err(WireCodecError::Kind);
338 }
339 }
340 Ok(())
341}
342
343fn borrowed(values: &[super::super::CallValue]) -> Vec<Value<'_>> {
344 values.iter().map(|value| Value { key: &value.key, payload: &value.payload }).collect()
345}
346
347fn owned(values: Vec<Value<'_>>) -> Result<Vec<super::super::CallValue>, WireCodecError> {
348 let bytes = values.iter().map(|value| value.key.len() + value.payload.len()).sum::<usize>();
349 if values.len() > super::super::PORTABLE_VALUES || bytes > super::super::PORTABLE_BYTES {
350 return Err(WireCodecError::Oversized);
351 }
352 Ok(values.into_iter().map(|value| super::super::CallValue { key: value.key.to_owned(), payload: value.payload.to_vec() }).collect())
353}
354
355fn encode(envelope: Envelope<'_>) -> Result<Vec<u8>, WireCodecError> {
356 let message = (VERSION, envelope);
357 let size = postcard::experimental::serialized_size(&message)?;
358 if size > LIMIT {
359 return Err(WireCodecError::Oversized);
360 }
361 Ok(postcard::to_allocvec(&message)?)
362}
363
364fn payload_size(envelope: Envelope<'_>, payload: usize) -> Result<usize, WireCodecError> {
365 let empty = postcard::experimental::serialized_size(&(VERSION, envelope))?;
366 let prefix = postcard::experimental::serialized_size(&payload)?;
367 add_size(add_size(empty, payload)?, prefix - 1)
368}
369
370fn add_size(left: usize, right: usize) -> Result<usize, WireCodecError> {
371 left.checked_add(right).ok_or(WireCodecError::Oversized)
372}
373
374fn decode(bytes: &[u8]) -> Result<Envelope<'_>, WireCodecError> {
375 if bytes.len() > LIMIT {
376 return Err(WireCodecError::Oversized);
377 }
378 let (version, rest) = postcard::take_from_bytes::<u8>(bytes)?;
379 if version != VERSION {
380 return Err(WireCodecError::Version);
381 }
382 let (envelope, rest) = postcard::take_from_bytes::<Envelope<'_>>(rest)?;
383 if !rest.is_empty() {
384 return Err(WireCodecError::Trailing);
385 }
386 Ok(envelope)
387}
388
389#[cfg(test)]
392mod tests {
393 use super::*;
394 use crate::operations::{CallValue, PORTABLE_BYTES, PORTABLE_VALUES};
395
396 #[test]
397 fn envelope_sizes_match() {
398 let sizes = [0, 1, 127, 128, 16383, 16384, 2097151, 2097152];
399 let ids = [0, 127, 128, u64::MAX];
400 for (index, size) in sizes.into_iter().enumerate() {
401 let operation = "λ".repeat([1, 63, 64, 65][index % 4]);
402 let invocation = ids[index % ids.len()];
403 let payload = vec![255; size];
404 let expected = WireCommand::prepare_encoded_len(invocation, &operation, size).unwrap();
405 let command = WireCommand::Prepare { invocation, operation, payload: payload.clone() };
406 assert_eq!(command.encode().unwrap().len(), expected);
407 let expected = WireAnswer::ready_encoded_len(size).unwrap();
408 assert_eq!(WireAnswer::Ready { payload }.encode().unwrap().len(), expected);
409 }
410 assert!(WireCommand::prepare_encoded_len(0, "op", usize::MAX).is_err());
411 assert!(WireAnswer::ready_encoded_len(usize::MAX).is_err());
412 }
413
414 #[test]
415 fn child_values_count() {
416 let mut maximal: Vec<_> =
417 (0..PORTABLE_VALUES).map(|index| CallValue { key: format!("k{index:02}"), payload: Vec::new() }).collect();
418 maximal[0].payload = vec![255; PORTABLE_BYTES - PORTABLE_VALUES * 3];
419 for values in [Vec::new(), vec![CallValue { key: "測定".into(), payload: vec![255; 128] }], maximal] {
420 for generation in [false, true] {
421 for submit in [false, true] {
422 let expected = WireChildCall::payload_encoded_len("op.λ", 16384, generation, &values).unwrap();
423 let call = WireChildCall {
424 context: [255; 32],
425 generation_context: generation.then_some([255; 32]),
426 values: values.clone(),
427 operation: "op.λ".into(),
428 payload: vec![255; 16384],
429 submit,
430 };
431 assert_eq!(call.encode().unwrap().len(), expected);
432 }
433 }
434 }
435 }
436
437 #[test]
438 fn child_validation_matches() {
439 let invalid = [
440 vec![CallValue { key: "key".into(), payload: vec![0; PORTABLE_BYTES] }],
441 vec![CallValue { key: "same".into(), payload: Vec::new() }; 2],
442 (0..=PORTABLE_VALUES).map(|index| CallValue { key: index.to_string(), payload: Vec::new() }).collect(),
443 ];
444 for values in invalid {
445 assert!(WireChildCall::payload_encoded_len("op", 0, false, &values).is_err());
446 let call = WireChildCall {
447 context: [0; 32],
448 generation_context: None,
449 values,
450 operation: "op".into(),
451 payload: Vec::new(),
452 submit: false,
453 };
454 assert!(call.encode().is_err());
455 }
456 assert!(WireChildCall::payload_encoded_len("op", usize::MAX, false, &[]).is_err());
457 }
458}