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 = 1;
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("unexpected operation envelope kind")]
62 Kind,
63
64 #[error("trailing operation envelope bytes")]
65 Trailing,
66
67 #[error("invalid operation envelope")]
68 Encoding(#[from] postcard::Error),
69}
70
71#[derive(Serialize, Deserialize)]
74enum Envelope<'a> {
75 Prepare {
76 invocation: u64,
77 operation: &'a str,
78 fingerprint: [u8; 32],
79 generation: u64,
80
81 #[serde(borrow, with = "payload")]
82 payload: &'a [u8],
83 },
84 Initialize {
85 invocation: u64,
86 context: [u8; 32],
87 },
88 Attempt {
89 invocation: u64,
90 context: [u8; 32],
91
92 #[serde(borrow)]
93 values: Vec<Value<'a>>,
94 },
95 Finalize {
96 invocation: u64,
97 },
98 Ready {
99 #[serde(borrow, with = "payload")]
100 payload: &'a [u8],
101 },
102 Failed {
103 failure: crate::Failure,
104 },
105 Child {
106 context: [u8; 32],
107 generation_context: Option<[u8; 32]>,
108
109 #[serde(borrow)]
110 values: Vec<Value<'a>>,
111
112 operation: &'a str,
113 fingerprint: [u8; 32],
114 generation: u64,
115
116 #[serde(borrow, with = "payload")]
117 payload: &'a [u8],
118
119 submit: bool,
120 },
121 PrepareGrpc {
122 invocation: u64,
123 operation: &'a str,
124 fingerprint: [u8; 32],
125 generation: u64,
126
127 #[serde(borrow, with = "payload")]
128 payload: &'a [u8],
129 },
130 StreamAttempt {
131 invocation: u64,
132 epoch: u64,
133 context: [u8; 32],
134
135 #[serde(borrow)]
136 values: Vec<Value<'a>>,
137 },
138 StreamSend {
139 invocation: u64,
140 epoch: u64,
141
142 #[serde(borrow, with = "payload")]
143 payload: &'a [u8],
144 },
145 StreamFinish {
146 invocation: u64,
147 epoch: u64,
148 },
149 StreamReceive {
150 invocation: u64,
151 epoch: u64,
152 },
153 StreamItem {
154 epoch: u64,
155 transferred: bool,
156
157 #[serde(borrow, with = "payload")]
158 payload: &'a [u8],
159 },
160 StreamEnd {
161 epoch: u64,
162 transferred: bool,
163 },
164 ChildControl {
165 context: [u8; 32],
166 instance: u64,
167 action: WireChildAction,
168 },
169}
170
171impl WireCommand {
176 pub fn prepare_encoded_len(
178 invocation: u64, operation: &str, fingerprint: [u8; 32], generation: u64, payload: usize,
179 ) -> Result<usize, WireCodecError> {
180 payload_size(
181 Envelope::Prepare {
182 invocation,
183 operation,
184 fingerprint,
185 generation,
186 payload: &[],
187 },
188 payload,
189 )
190 }
191
192 pub fn encode(&self) -> Result<Vec<u8>, WireCodecError> {
193 encode(self.envelope())
194 }
195
196 pub fn decode(bytes: &[u8]) -> Result<Self, WireCodecError> {
197 Ok(match decode(bytes)? {
198 Envelope::Prepare {
199 invocation,
200 operation,
201 fingerprint,
202 generation,
203 payload,
204 } => Self::Prepare {
205 invocation,
206 operation: operation.to_owned(),
207 fingerprint,
208 generation,
209 payload: payload.to_vec(),
210 },
211 Envelope::PrepareGrpc {
212 invocation,
213 operation,
214 fingerprint,
215 generation,
216 payload,
217 } => Self::PrepareGrpc {
218 invocation,
219 operation: operation.to_owned(),
220 fingerprint,
221 generation,
222 payload: payload.to_vec(),
223 },
224 Envelope::Initialize { invocation, context } => Self::Initialize { invocation, context },
225 Envelope::Attempt { invocation, context, values } => Self::Attempt {
226 invocation,
227 context,
228 values: owned(values)?,
229 },
230 Envelope::Finalize { invocation } => Self::Finalize { invocation },
231 Envelope::StreamAttempt { invocation, epoch, context, values } => Self::StreamAttempt {
232 invocation,
233 epoch,
234 context,
235 values: owned(values)?,
236 },
237 Envelope::StreamSend { invocation, epoch, payload } => Self::StreamSend {
238 invocation,
239 epoch,
240 payload: payload.to_vec(),
241 },
242 Envelope::StreamFinish { invocation, epoch } => Self::StreamFinish { invocation, epoch },
243 Envelope::StreamReceive { invocation, epoch } => Self::StreamReceive { invocation, epoch },
244 _ => return Err(WireCodecError::Kind),
245 })
246 }
247
248 fn envelope(&self) -> Envelope<'_> {
249 match self {
250 Self::Prepare {
251 invocation,
252 operation,
253 fingerprint,
254 generation,
255 payload,
256 } => Envelope::Prepare {
257 invocation: *invocation,
258 operation,
259 fingerprint: *fingerprint,
260 generation: *generation,
261 payload,
262 },
263 Self::PrepareGrpc {
264 invocation,
265 operation,
266 fingerprint,
267 generation,
268 payload,
269 } => Envelope::PrepareGrpc {
270 invocation: *invocation,
271 operation,
272 fingerprint: *fingerprint,
273 generation: *generation,
274 payload,
275 },
276 Self::Initialize { invocation, context } => Envelope::Initialize {
277 invocation: *invocation,
278 context: *context,
279 },
280 Self::Attempt { invocation, context, values } => Envelope::Attempt {
281 invocation: *invocation,
282 context: *context,
283 values: borrowed(values),
284 },
285 Self::StreamAttempt { invocation, epoch, context, values } => Envelope::StreamAttempt {
286 invocation: *invocation,
287 epoch: *epoch,
288 context: *context,
289 values: borrowed(values),
290 },
291 Self::StreamSend { invocation, epoch, payload } => Envelope::StreamSend {
292 invocation: *invocation,
293 epoch: *epoch,
294 payload,
295 },
296 Self::StreamFinish { invocation, epoch } => Envelope::StreamFinish { invocation: *invocation, epoch: *epoch },
297 Self::StreamReceive { invocation, epoch } => Envelope::StreamReceive { invocation: *invocation, epoch: *epoch },
298 Self::Finalize { invocation } => Envelope::Finalize { invocation: *invocation },
299 }
300 }
301}
302
303impl WireAnswer {
304 pub fn ready_encoded_len(payload: usize) -> Result<usize, WireCodecError> {
306 payload_size(Envelope::Ready { payload: &[] }, payload)
307 }
308
309 pub fn encode(&self) -> Result<Vec<u8>, WireCodecError> {
310 let envelope = match self {
311 Self::Ready { payload } => Envelope::Ready { payload },
312 Self::Stream {
313 epoch,
314 transferred,
315 payload: Some(payload),
316 } => Envelope::StreamItem {
317 epoch: *epoch,
318 transferred: *transferred,
319 payload,
320 },
321 Self::Stream { epoch, transferred, payload: None } => Envelope::StreamEnd { epoch: *epoch, transferred: *transferred },
322 Self::Failed { failure } => Envelope::Failed { failure: failure.clone() },
323 };
324 encode(envelope)
325 }
326
327 pub fn decode(bytes: &[u8]) -> Result<Self, WireCodecError> {
328 Ok(match decode(bytes)? {
329 Envelope::StreamItem { epoch, transferred, payload } => Self::Stream {
330 epoch,
331 transferred,
332 payload: Some(payload.to_vec()),
333 },
334 Envelope::StreamEnd { epoch, transferred } => Self::Stream { epoch, transferred, payload: None },
335 Envelope::Ready { payload } => Self::Ready { payload: payload.to_vec() },
336 Envelope::Failed { failure } => Self::Failed { failure },
337 _ => return Err(WireCodecError::Kind),
338 })
339 }
340}
341
342impl WireChildCall {
343 pub fn payload_encoded_len(
346 operation: &str, generation: u64, payload: usize, generation_context: bool, values: &[super::super::CallValue],
347 ) -> Result<usize, WireCodecError> {
348 check_values(values)?;
349 let envelope = Envelope::Child {
350 context: [0; 32],
351 generation_context: generation_context.then_some([0; 32]),
352 values: Vec::new(),
353 operation,
354 fingerprint: [0; 32],
355 generation,
356 payload: &[],
357 submit: false,
358 };
359 let count = postcard::experimental::serialized_size(&values.len())?;
360 let mut size = add_size(payload_size(envelope, payload)?, count - 1)?;
361 for value in values {
362 let value = Value { key: &value.key, payload: &value.payload };
363 size = add_size(size, postcard::experimental::serialized_size(&value)?)?;
364 }
365 Ok(size)
366 }
367
368 pub fn encode(&self) -> Result<Vec<u8>, WireCodecError> {
369 check_values(&self.values)?;
370 encode(Envelope::Child {
371 context: self.context,
372 generation_context: self.generation_context,
373 values: borrowed(&self.values),
374 operation: &self.operation,
375 fingerprint: self.fingerprint,
376 generation: self.generation,
377 payload: &self.payload,
378 submit: self.submit,
379 })
380 }
381
382 pub fn decode(bytes: &[u8]) -> Result<Self, WireCodecError> {
383 match WireChild::decode(bytes)? {
384 WireChild::Call(call) => Ok(call),
385 WireChild::Control(_) => Err(WireCodecError::Kind),
386 }
387 }
388}
389
390impl WireChildControl {
391 pub fn encode(&self) -> Result<Vec<u8>, WireCodecError> {
392 encode(Envelope::ChildControl {
393 context: self.context,
394 instance: self.instance,
395 action: self.action,
396 })
397 }
398}
399
400impl WireChild {
401 pub fn decode(bytes: &[u8]) -> Result<Self, WireCodecError> {
402 let child = match decode(bytes)? {
403 Envelope::Child {
404 context,
405 generation_context,
406 values,
407 operation,
408 fingerprint,
409 generation,
410 payload,
411 submit,
412 } => Self::Call(WireChildCall {
413 context,
414 generation_context,
415 values: owned(values)?,
416 operation: operation.to_owned(),
417 fingerprint,
418 generation,
419 payload: payload.to_vec(),
420 submit,
421 }),
422 Envelope::ChildControl { context, instance, action } => Self::Control(WireChildControl { context, instance, action }),
423 _ => return Err(WireCodecError::Kind),
424 };
425 Ok(child)
426 }
427
428 pub fn context(&self) -> &[u8; 32] {
429 match self {
430 Self::Call(call) => &call.context,
431 Self::Control(control) => &control.context,
432 }
433 }
434}
435
436fn check_values(values: &[super::super::CallValue]) -> Result<(), WireCodecError> {
441 let bytes = values.iter().map(|value| value.key.len().saturating_add(value.payload.len())).fold(0usize, usize::saturating_add);
442 if values.len() > super::super::PORTABLE_VALUES || bytes > super::super::PORTABLE_BYTES {
443 return Err(WireCodecError::Oversized);
444 }
445 for (index, value) in values.iter().enumerate() {
446 if values[..index].iter().any(|previous| previous.key == value.key) {
447 return Err(WireCodecError::Kind);
448 }
449 }
450 Ok(())
451}
452
453fn borrowed(values: &[super::super::CallValue]) -> Vec<Value<'_>> {
454 values.iter().map(|value| Value { key: &value.key, payload: &value.payload }).collect()
455}
456
457fn owned(values: Vec<Value<'_>>) -> Result<Vec<super::super::CallValue>, WireCodecError> {
458 let bytes = values.iter().map(|value| value.key.len() + value.payload.len()).sum::<usize>();
459 if values.len() > super::super::PORTABLE_VALUES || bytes > super::super::PORTABLE_BYTES {
460 return Err(WireCodecError::Oversized);
461 }
462 Ok(values
463 .into_iter()
464 .map(|value| super::super::CallValue {
465 key: value.key.to_owned(),
466 payload: value.payload.to_vec(),
467 })
468 .collect())
469}
470
471fn encode(envelope: Envelope<'_>) -> Result<Vec<u8>, WireCodecError> {
472 let message = (VERSION, envelope);
473 let size = postcard::experimental::serialized_size(&message)?;
474 if size > LIMIT {
475 return Err(WireCodecError::Oversized);
476 }
477 Ok(postcard::to_allocvec(&message)?)
478}
479
480fn payload_size(envelope: Envelope<'_>, payload: usize) -> Result<usize, WireCodecError> {
481 let empty = postcard::experimental::serialized_size(&(VERSION, envelope))?;
482 let prefix = postcard::experimental::serialized_size(&payload)?;
483 add_size(add_size(empty, payload)?, prefix - 1)
484}
485
486fn add_size(left: usize, right: usize) -> Result<usize, WireCodecError> {
487 left.checked_add(right).ok_or(WireCodecError::Oversized)
488}
489
490fn decode(bytes: &[u8]) -> Result<Envelope<'_>, WireCodecError> {
491 if bytes.len() > LIMIT {
492 return Err(WireCodecError::Oversized);
493 }
494 let (_version, rest) = postcard::take_from_bytes::<u8>(bytes)?;
495 let (envelope, rest) = postcard::take_from_bytes::<Envelope<'_>>(rest)?;
496 if !rest.is_empty() {
497 return Err(WireCodecError::Trailing);
498 }
499 Ok(envelope)
500}
501
502#[cfg(test)]
505mod tests {
506 use super::*;
507 use crate::operations::{CallValue, PORTABLE_BYTES, PORTABLE_VALUES};
508
509 #[test]
510 fn envelope_sizes_match() {
511 let sizes = [0, 1, 127, 128, 16383, 16384, 2097151, 2097152];
512 let ids = [0, 127, 128, u64::MAX];
513 for (index, size) in sizes.into_iter().enumerate() {
514 let operation = "λ".repeat([1, 63, 64, 65][index % 4]);
515 let invocation = ids[index % ids.len()];
516 let payload = vec![255; size];
517 let expected = WireCommand::prepare_encoded_len(invocation, &operation, [255; 32], invocation, size).unwrap();
518 let command = WireCommand::Prepare {
519 invocation,
520 operation,
521 fingerprint: [255; 32],
522 generation: invocation,
523 payload: payload.clone(),
524 };
525 assert_eq!(command.encode().unwrap().len(), expected);
526 let expected = WireAnswer::ready_encoded_len(size).unwrap();
527 assert_eq!(WireAnswer::Ready { payload }.encode().unwrap().len(), expected);
528 }
529 assert!(WireCommand::prepare_encoded_len(0, "op", [255; 32], u64::MAX, usize::MAX).is_err());
530 assert!(WireAnswer::ready_encoded_len(usize::MAX).is_err());
531 }
532
533 #[test]
534 fn prepare_identity_roundtrips() {
535 let command = WireCommand::Prepare {
536 invocation: 127,
537 operation: "fixture.read".into(),
538 fingerprint: [173; 32],
539 generation: u64::MAX,
540 payload: vec![1, 2],
541 };
542 let bytes = command.encode().unwrap();
543 assert!(matches!(WireCommand::decode(&bytes).unwrap(), WireCommand::Prepare {
544 invocation: 127, operation, fingerprint, generation: u64::MAX, payload
545 } if operation == "fixture.read" && fingerprint == [173; 32] && payload == [1, 2]));
546 }
547
548 #[test]
549 fn child_values_count() {
550 let mut maximal: Vec<_> = (0..PORTABLE_VALUES)
551 .map(|index| CallValue {
552 key: format!("k{index:02}"),
553 payload: Vec::new(),
554 })
555 .collect();
556 maximal[0].payload = vec![255; PORTABLE_BYTES - PORTABLE_VALUES * 3];
557 for values in [
558 Vec::new(),
559 vec![CallValue {
560 key: "測定".into(),
561 payload: vec![255; 128],
562 }],
563 maximal,
564 ] {
565 for generation in [false, true] {
566 for submit in [false, true] {
567 let expected = WireChildCall::payload_encoded_len("op.λ", u64::MAX, 16384, generation, &values).unwrap();
568 let call = WireChildCall {
569 context: [255; 32],
570 generation_context: generation.then_some([255; 32]),
571 values: values.clone(),
572 operation: "op.λ".into(),
573 fingerprint: [1; 32],
574 generation: u64::MAX,
575 payload: vec![255; 16384],
576 submit,
577 };
578 assert_eq!(call.encode().unwrap().len(), expected);
579 }
580 }
581 }
582 }
583
584 #[test]
585 fn child_validation_matches() {
586 let invalid = [
587 vec![CallValue {
588 key: "key".into(),
589 payload: vec![0; PORTABLE_BYTES],
590 }],
591 vec![CallValue { key: "same".into(), payload: Vec::new() }; 2],
592 (0..=PORTABLE_VALUES)
593 .map(|index| CallValue {
594 key: index.to_string(),
595 payload: Vec::new(),
596 })
597 .collect(),
598 ];
599 for values in invalid {
600 assert!(WireChildCall::payload_encoded_len("op", u64::MAX, 0, false, &values).is_err());
601 let call = WireChildCall {
602 context: [0; 32],
603 generation_context: None,
604 values,
605 operation: "op".into(),
606 fingerprint: [1; 32],
607 generation: u64::MAX,
608 payload: Vec::new(),
609 submit: false,
610 };
611 assert!(call.encode().is_err());
612 }
613 assert!(WireChildCall::payload_encoded_len("op", u64::MAX, usize::MAX, false, &[]).is_err());
614 }
615}