Skip to main content

serde_multipart/
lib.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! Serde codec for multipart messages.
10//!
11//! Using [`serialize`] / [`deserialize`], fields typed [`Part`] are extracted
12//! from the main payload and appended to a list of `parts`. Each part is backed by
13//! [`Vec<bytes::Bytes>`] for cheap, zero-copy sharing.
14//!
15//! On decode, the body and its parts are reassembled into the original value
16//! without copying.
17//!
18//! The on-the-wire form is a [`Message`] (body + parts). Your transport sends
19//! and receives [`Message`]s; the codec reconstructs the value, enabling
20//! efficient network I/O without compacting data into a single buffer.
21//!
22//! Implementation note: this crate uses Rust's min_specialization feature to enable
23//! the use of [`Part`]s with any Serde serializer or deserializer. This feature
24//! is fairly restrictive, and thus the API offered by [`serialize`] / [`deserialize`]
25//! is not customizable. If customization is needed, you need to add specialization
26//! implementations for these codecs. See [`part::PartSerializer`] and [`part::PartDeserializer`]
27//! for details.
28
29#![feature(min_specialization)]
30
31use std::cell::UnsafeCell;
32use std::cmp::min;
33use std::collections::VecDeque;
34use std::io::IoSlice;
35use std::ptr::NonNull;
36
37use bincode::Options;
38use bytes::Buf;
39use bytes::BufMut;
40use bytes::buf::UninitSlice;
41
42mod codec;
43mod de;
44mod part;
45mod ser;
46use bytes::Bytes;
47use bytes::BytesMut;
48pub use codec::PartCodec;
49pub use codec::deserialize_part_codec;
50pub use codec::serialize_part_codec;
51pub use part::Part;
52use serde::Deserialize;
53use serde::Serialize;
54use serde::de::DeserializeOwned;
55use typeuri::Named;
56
57/// The type of error returned by typed part operations.
58#[derive(Debug, thiserror::Error)]
59pub enum Error {
60    /// Errors returned from bincode.
61    #[error(transparent)]
62    Bincode(#[from] bincode::Error),
63
64    /// Type mismatch during deserialization.
65    #[error("type mismatch: expected {expected}, found {actual}")]
66    TypeMismatch {
67        expected: &'static str,
68        actual: String,
69    },
70
71    /// Errors returned from part codec conversions.
72    #[error("codec error: {0}")]
73    Codec(String),
74}
75
76/// A specialized result type for typed part operations.
77pub type Result<T, E = Error> = std::result::Result<T, E>;
78
79const FRAME_TYPEHASH_FLAG: u64 = 1 << 63;
80const FRAME_LEN_MASK: u64 = !FRAME_TYPEHASH_FLAG;
81
82/// A multi-part message, comprising a message body and a list of parts.
83/// Messages only contain references to underlying byte buffers and are
84/// cheaply cloned.
85#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
86pub struct Message {
87    body: Part,
88    parts: Vec<Part>,
89}
90
91impl Message {
92    /// Returns a new message with the given body and parts.
93    pub fn from_body_and_parts(body: Part, parts: Vec<Part>) -> Self {
94        Self { body, parts }
95    }
96
97    /// The body of the message.
98    pub fn body(&self) -> &Part {
99        &self.body
100    }
101
102    /// The list of parts of the message.
103    pub fn parts(&self) -> &[Part] {
104        &self.parts
105    }
106
107    /// Returns the total number of parts (excluding the body) in the message.
108    pub fn num_parts(&self) -> usize {
109        self.parts.len()
110    }
111
112    /// Returns the total size (in bytes) of the message.
113    pub fn len(&self) -> usize {
114        self.body.len() + self.parts.iter().map(|part| part.len()).sum::<usize>()
115    }
116
117    /// Returns whether the message is empty. It is always false, since the body
118    /// is always defined.
119    pub fn is_empty(&self) -> bool {
120        self.body.is_empty() && self.parts.iter().all(|part| part.is_empty())
121    }
122
123    /// Convert this message into its constituent components.
124    pub fn into_inner(self) -> (Part, Vec<Part>) {
125        (self.body, self.parts)
126    }
127
128    /// Visit every typed part containing a `T`-typed value.
129    pub fn visit_parts_mut<T, E>(
130        &mut self,
131        mut f: impl FnMut(&mut T) -> std::result::Result<(), E>,
132    ) -> std::result::Result<(), E>
133    where
134        T: Serialize + DeserializeOwned + Named,
135        E: From<Error>,
136    {
137        for part in &mut self.parts {
138            if part.is::<T>() {
139                let mut value = part.deserialized::<T>().map_err(E::from)?;
140                f(&mut value)?;
141                *part = Part::serialize(&value).map_err(E::from)?;
142            }
143        }
144        Ok(())
145    }
146
147    /// Returns the total size (in bytes) of the message when it is framed.
148    pub fn frame_len(&self) -> usize {
149        Self::part_frame_len(&self.body)
150            + self.parts.iter().map(Self::part_frame_len).sum::<usize>()
151    }
152
153    /// Efficiently frames a message containing the body and all of its parts
154    /// using a simple frame-length encoding:
155    ///
156    /// ```text
157    /// +--------------------+-------------------+--------------------+-------------------+   ...   +
158    /// | body_tag (u64 BE)  |   body bytes      | part1_tag (u64 BE) |   part1 bytes     |         |
159    /// +--------------------+-------------------+--------------------+-------------------+         +
160    ///                                                                                      repeat
161    ///                                                                                        for
162    ///                                                                                      each part
163    /// ```
164    ///
165    /// The high bit of each tag indicates whether the part is typed. If set,
166    /// the lower 63 bits hold the part length, and the tag is followed by the
167    /// part typehash as a `u64 BE` before the part bytes.
168    pub fn framed(self) -> Frame {
169        let (body, parts) = self.into_inner();
170
171        let mut buffers = Vec::with_capacity(
172            Self::part_frame_buffers(&body)
173                + parts.iter().map(Self::part_frame_buffers).sum::<usize>(),
174        );
175
176        Self::push_framed_part(body, &mut buffers);
177
178        for part in parts {
179            Self::push_framed_part(part, &mut buffers);
180        }
181
182        Frame::from_buffers(buffers)
183    }
184
185    /// Reassembles a message from a framed encoding.
186    pub fn from_framed(mut buf: Bytes) -> Result<Self, std::io::Error> {
187        if buf.len() < 8 {
188            return Err(std::io::ErrorKind::UnexpectedEof.into());
189        }
190        let body = Self::split_part(&mut buf)?;
191        let mut parts = Vec::new();
192        while !buf.is_empty() {
193            parts.push(Self::split_part(&mut buf)?);
194        }
195        Ok(Self { body, parts })
196    }
197
198    fn part_frame_len(part: &Part) -> usize {
199        8usize
200            .checked_add(part.typehash().map_or(0, |_| 8))
201            .and_then(|len| len.checked_add(part.len()))
202            .expect("part frame length exceeds usize")
203    }
204
205    fn part_frame_buffers(part: &Part) -> usize {
206        1 + part.typehash().map_or(0, |_| 1) + part.num_fragments()
207    }
208
209    fn push_framed_part(part: Part, buffers: &mut Vec<Bytes>) {
210        let typehash = part.typehash();
211        let tag = Self::frame_tag(part.len(), typehash);
212        buffers.push(Bytes::from_owner(tag.to_be_bytes()));
213        if let Some(typehash) = typehash {
214            buffers.push(Bytes::from_owner(typehash.to_be_bytes()));
215        }
216        for fragment in part.into_fragments() {
217            buffers.push(fragment);
218        }
219    }
220
221    fn frame_tag(len: usize, typehash: Option<u64>) -> u64 {
222        let len = u64::try_from(len).expect("part length exceeds u64");
223        if len > FRAME_LEN_MASK {
224            panic!("part length exceeds 63-bit frame limit");
225        }
226        len | typehash.map_or(0, |_| FRAME_TYPEHASH_FLAG)
227    }
228
229    fn frame_len_from_tag(tag: u64) -> Result<usize, std::io::Error> {
230        usize::try_from(tag & FRAME_LEN_MASK).map_err(|_| {
231            std::io::Error::new(
232                std::io::ErrorKind::InvalidData,
233                "part length exceeds addressable memory",
234            )
235        })
236    }
237
238    fn split_part(buf: &mut Bytes) -> Result<Part, std::io::Error> {
239        if buf.len() < 8 {
240            return Err(std::io::ErrorKind::UnexpectedEof.into());
241        }
242        let tag = buf.get_u64();
243        let typehash = if tag & FRAME_TYPEHASH_FLAG == 0 {
244            None
245        } else {
246            if buf.len() < 8 {
247                return Err(std::io::ErrorKind::UnexpectedEof.into());
248            }
249            Some(buf.get_u64())
250        };
251        let at = Self::frame_len_from_tag(tag)?;
252        if buf.len() < at {
253            return Err(std::io::ErrorKind::UnexpectedEof.into());
254        }
255        Ok(Part::from_typehash_and_fragments(
256            typehash,
257            vec![buf.split_to(at)],
258        ))
259    }
260}
261
262/// An encoded [`Message`] frame. Implements [`bytes::Buf`],
263/// and supports vectored writes. Thus, `Frame` is like a reader
264/// of an encoded [`Message`].
265#[derive(Clone)]
266pub struct Frame {
267    buffers: VecDeque<Bytes>,
268}
269
270impl Frame {
271    /// Construct a new frame from the provided buffers. The frame is a
272    /// concatenation of these buffers.
273    fn from_buffers(buffers: Vec<Bytes>) -> Self {
274        let mut buffers: VecDeque<Bytes> = buffers.into();
275        buffers.retain(|buf| !buf.is_empty());
276        Self { buffers }
277    }
278
279    /// **DO NOT USE THIS**
280    pub fn illegal_unipart_frame(body: Bytes) -> Self {
281        Self {
282            buffers: vec![body].into(),
283        }
284    }
285}
286
287impl Buf for Frame {
288    fn remaining(&self) -> usize {
289        self.buffers.iter().map(|buf| buf.remaining()).sum()
290    }
291
292    fn chunk(&self) -> &[u8] {
293        match self.buffers.front() {
294            Some(buf) => buf.chunk(),
295            None => &[],
296        }
297    }
298
299    fn advance(&mut self, mut cnt: usize) {
300        while cnt > 0 {
301            let Some(buf) = self.buffers.front_mut() else {
302                panic!("advanced beyond the buffer size");
303            };
304
305            if cnt >= buf.remaining() {
306                cnt -= buf.remaining();
307                self.buffers.pop_front();
308                continue;
309            }
310
311            buf.advance(cnt);
312            cnt = 0;
313        }
314    }
315
316    // We implement our own chunks_vectored here, as the default implementation
317    // does not do any vectoring (returning only a single IoSlice at a time).
318    fn chunks_vectored<'a>(&'a self, dst: &mut [IoSlice<'a>]) -> usize {
319        let n = min(dst.len(), self.buffers.len());
320        for (i, slot) in dst.iter_mut().enumerate().take(n) {
321            *slot = IoSlice::new(self.buffers[i].chunk());
322        }
323        n
324    }
325}
326
327/// An unsafe cell of a [`BytesMut`]. This is used to implement an io::Writer
328/// for the serializer without exposing lifetime parameters (which cannot be)
329/// specialized.
330struct UnsafeBufCell {
331    buf: UnsafeCell<BytesMut>,
332}
333
334impl UnsafeBufCell {
335    /// Create a new cell from a [`BytesMut`].
336    fn from_bytes_mut(bytes: BytesMut) -> Self {
337        Self {
338            buf: UnsafeCell::new(bytes),
339        }
340    }
341
342    /// Convert this cell into its underlying [`BytesMut`].
343    fn into_inner(self) -> BytesMut {
344        self.buf.into_inner()
345    }
346
347    /// Borrow the cell, without lifetime checks. The caller must guarantee that
348    /// the returned cell cannot be used after the cell is dropped (usually through
349    /// [`UnsafeBufCell::into_inner`]).
350    unsafe fn borrow_unchecked(&self) -> UnsafeBufCellRef {
351        let ptr =
352            // SAFETY: the user is providing the necessary invariants
353            unsafe { NonNull::new_unchecked(self.buf.get()) };
354        UnsafeBufCellRef { ptr }
355    }
356}
357
358/// A borrowed reference to an [`UnsafeBufCell`].
359struct UnsafeBufCellRef {
360    ptr: NonNull<BytesMut>,
361}
362
363/// SAFETY: we're extending the implementation of the underlying [`BytesMut`];
364/// adding an additional layer of danger by disregarding lifetimes.
365unsafe impl BufMut for UnsafeBufCellRef {
366    fn remaining_mut(&self) -> usize {
367        // SAFETY: extending the implementation of the underlying [`BytesMut`]
368        unsafe { self.ptr.as_ref().remaining_mut() }
369    }
370
371    unsafe fn advance_mut(&mut self, cnt: usize) {
372        // SAFETY: extending the implementation of the underlying [`BytesMut`]
373        unsafe { self.ptr.as_mut().advance_mut(cnt) }
374    }
375
376    fn chunk_mut(&mut self) -> &mut UninitSlice {
377        // SAFETY: extending the implementation of the underlying [`BytesMut`]
378        unsafe { self.ptr.as_mut().chunk_mut() }
379    }
380}
381
382/// Serialize the provided value into a multipart message. The value is encoded using an
383/// extended version of [`bincode`] that skips serializing [`Part`]s, which are instead
384/// held directly by the returned message.
385///
386/// Serialize uses the same codec options as [`bincode::serialize`] / [`bincode::deserialize`].
387/// These are currently not customizable unless an explicit specialization is also provided.
388pub fn serialize_bincode<S: ?Sized + serde::Serialize>(
389    value: &S,
390) -> Result<Message, bincode::Error> {
391    let buffer = UnsafeBufCell::from_bytes_mut(BytesMut::new());
392    // SAFETY: we know here that, once the below "value.serialize()" is done, there are no more
393    // extant references to this buffer; we are thus safe to reclaim the buffer into the message
394    let buffer_borrow = unsafe { buffer.borrow_unchecked() };
395    let mut serializer: part::BincodeSerializer =
396        ser::bincode::Serializer::new(bincode::Serializer::new(buffer_borrow.writer(), options()));
397    value.serialize(&mut serializer)?;
398    Ok(Message {
399        body: Part::from_fragments(vec![buffer.into_inner().freeze()]),
400        parts: serializer.into_parts(),
401    })
402}
403
404/// Deserialize a message serialized by `[serialize]`, stitching together the original
405/// message without copying the underlying buffers.
406pub fn deserialize_bincode<T>(message: Message) -> Result<T, bincode::Error>
407where
408    T: serde::de::DeserializeOwned,
409{
410    let (body, parts) = message.into_inner();
411    let mut deserializer = part::BincodeDeserializer::new(
412        bincode::Deserializer::with_reader(body.into_bytes().reader(), options()),
413        parts.into(),
414    );
415    let value = T::deserialize(&mut deserializer)?;
416    // Check that all parts were consumed:
417    deserializer.end()?;
418    Ok(value)
419}
420
421/// Construct the set of options used by the specialized serializer and deserializer.
422fn options() -> part::BincodeOptionsType {
423    bincode::DefaultOptions::new()
424        .with_fixint_encoding()
425        .allow_trailing_bytes()
426}
427
428#[cfg(test)]
429mod tests {
430    use std::assert_matches;
431    use std::collections::HashMap;
432    use std::net::SocketAddr;
433    use std::net::SocketAddrV6;
434
435    use bytes::BufMut;
436    use proptest::prelude::*;
437    use proptest_derive::Arbitrary;
438    use serde::Deserialize;
439    use serde::Serialize;
440    use serde::de::DeserializeOwned;
441    use typeuri::Named;
442
443    use super::*;
444
445    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Named)]
446    struct TypedPayload {
447        label: String,
448        value: u64,
449    }
450
451    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Named)]
452    struct TypedPayloadRepr {
453        label: String,
454        value: u64,
455    }
456
457    #[derive(Debug, Clone, PartialEq, Eq, Named)]
458    struct CodecPayload {
459        label: String,
460        value: u64,
461    }
462
463    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Named)]
464    struct CodecPayloadRepr {
465        label: String,
466        value: u64,
467    }
468
469    impl PartCodec for CodecPayload {
470        type Repr = CodecPayloadRepr;
471    }
472
473    impl TryFrom<&CodecPayload> for CodecPayloadRepr {
474        type Error = Error;
475
476        fn try_from(value: &CodecPayload) -> Result<Self> {
477            Ok(Self {
478                label: value.label.clone(),
479                value: value.value,
480            })
481        }
482    }
483
484    impl TryFrom<CodecPayloadRepr> for CodecPayload {
485        type Error = Error;
486
487        fn try_from(repr: CodecPayloadRepr) -> Result<Self> {
488            Ok(Self {
489                label: repr.label,
490                value: repr.value,
491            })
492        }
493    }
494
495    impl Serialize for CodecPayload {
496        fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
497        where
498            S: serde::Serializer,
499        {
500            serialize_part_codec(self, serializer)
501        }
502    }
503
504    impl<'de> Deserialize<'de> for CodecPayload {
505        fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
506        where
507            D: serde::Deserializer<'de>,
508        {
509            deserialize_part_codec(deserializer)
510        }
511    }
512
513    #[derive(Debug, Clone, PartialEq, Eq, Named)]
514    struct MacroCodecPayload {
515        label: String,
516        value: u64,
517    }
518
519    #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Named)]
520    struct MacroCodecPayloadRepr {
521        label: String,
522        value: u64,
523    }
524
525    crate::part_codec! {
526        impl MacroCodecPayload
527        {
528            type Repr = MacroCodecPayloadRepr;
529        }
530    }
531
532    impl TryFrom<&MacroCodecPayload> for MacroCodecPayloadRepr {
533        type Error = Error;
534
535        fn try_from(value: &MacroCodecPayload) -> Result<Self> {
536            Ok(Self {
537                label: value.label.clone(),
538                value: value.value,
539            })
540        }
541    }
542
543    impl TryFrom<MacroCodecPayloadRepr> for MacroCodecPayload {
544        type Error = Error;
545
546        fn try_from(repr: MacroCodecPayloadRepr) -> Result<Self> {
547            Ok(Self {
548                label: repr.label,
549                value: repr.value,
550            })
551        }
552    }
553
554    fn test_roundtrip<T>(value: T, expected_parts: usize)
555    where
556        T: Serialize + DeserializeOwned + PartialEq + std::fmt::Debug,
557    {
558        // Test plain serialization roundtrip:
559        let message = serialize_bincode(&value).unwrap();
560        assert_eq!(message.num_parts(), expected_parts);
561        let deserialized_value = deserialize_bincode(message.clone()).unwrap();
562        assert_eq!(value, deserialized_value);
563
564        // Framing roundtrip:
565        let mut framed = message.clone().framed();
566        let framed = framed.copy_to_bytes(framed.remaining());
567        let unframed_message = Message::from_framed(framed).unwrap();
568        assert_eq!(message, unframed_message);
569
570        // Bincode passthrough:
571        let bincode_serialized = bincode::serialize(&value).unwrap();
572        let bincode_deserialized = bincode::deserialize(&bincode_serialized).unwrap();
573        assert_eq!(value, bincode_deserialized);
574    }
575
576    #[test]
577    fn test_specialized_serializer_basic() {
578        test_roundtrip(Part::from("hello"), 1);
579    }
580
581    #[test]
582    fn test_specialized_serializer_compound() {
583        test_roundtrip(vec![Part::from("hello"), Part::from("world")], 2);
584        test_roundtrip((Part::from("hello"), 1, 2, 3, Part::from("world")), 2);
585        test_roundtrip(
586            {
587                #[derive(Serialize, Deserialize, Debug, PartialEq)]
588                struct U {
589                    parts: Vec<Part>,
590                }
591                #[derive(Serialize, Deserialize, Debug, PartialEq)]
592                enum E {
593                    First(Part),
594                    Second(String),
595                }
596
597                #[derive(Serialize, Deserialize, Debug, PartialEq)]
598                struct T {
599                    field2: String,
600                    field3: Part,
601                    field4: Part,
602                    field5: Vec<U>,
603                    field6: E,
604                }
605
606                T {
607                    field2: "hello".to_string(),
608                    field3: Part::from("hello"),
609                    field4: Part::from("world"),
610                    field5: vec![
611                        U {
612                            parts: vec![Part::from("hello"), Part::from("world")],
613                        },
614                        U {
615                            parts: vec![Part::from("five"), Part::from("six"), Part::from("seven")],
616                        },
617                    ],
618                    field6: E::First(Part::from("eight")),
619                }
620            },
621            8,
622        );
623        test_roundtrip(
624            {
625                #[derive(Serialize, Deserialize, Debug, PartialEq)]
626                struct T {
627                    field1: u64,
628                    field2: String,
629                    field3: Part,
630                    field4: Part,
631                    field5: u64,
632                }
633                T {
634                    field1: 1,
635                    field2: "hello".to_string(),
636                    field3: Part::from("hello"),
637                    field4: Part::from("world"),
638                    field5: 2,
639                }
640            },
641            2,
642        );
643    }
644
645    #[test]
646    fn test_recursive_message() {
647        let message = serialize_bincode(&[Part::from("hello"), Part::from("world")]).unwrap();
648        let message_message = serialize_bincode(&message).unwrap();
649
650        // message.body + message.parts (x2):
651        assert_eq!(message_message.num_parts(), 3);
652    }
653
654    #[test]
655    fn test_typed_part() {
656        let value = TypedPayload {
657            label: "hello".to_string(),
658            value: 42,
659        };
660        let part = Part::serialize(&value).unwrap();
661
662        assert_eq!(part.typehash(), Some(TypedPayload::typehash()));
663        assert!(part.is::<TypedPayload>());
664        assert_eq!(part.deserialized::<TypedPayload>().unwrap(), value);
665
666        let err = part.deserialized::<String>().unwrap_err();
667        assert_matches!(
668            err,
669            Error::TypeMismatch { expected, actual }
670                if expected == String::typename() && actual == TypedPayload::typehash().to_string()
671        );
672
673        let untyped = Part::from("hello");
674        let err = untyped.deserialized::<TypedPayload>().unwrap_err();
675        assert_matches!(
676            err,
677            Error::TypeMismatch { actual, .. } if actual == "unknown"
678        );
679
680        let bincode_serialized = bincode::serialize(&part).unwrap();
681        let bincode_deserialized: Part = bincode::deserialize(&bincode_serialized).unwrap();
682        assert_eq!(part, bincode_deserialized);
683        assert_eq!(
684            bincode_deserialized.deserialized::<TypedPayload>().unwrap(),
685            value
686        );
687    }
688
689    #[test]
690    fn test_typed_part_with_distinct_repr() {
691        let repr = TypedPayloadRepr {
692            label: "repr".to_string(),
693            value: 99,
694        };
695        let part = Part::serialize_as::<TypedPayload, _>(&repr).unwrap();
696
697        assert_eq!(part.typehash(), Some(TypedPayload::typehash()));
698        assert!(part.is::<TypedPayload>());
699        assert_eq!(
700            part.deserialized_as::<TypedPayload, TypedPayloadRepr>()
701                .unwrap(),
702            repr
703        );
704
705        let err = part
706            .deserialized_as::<String, TypedPayloadRepr>()
707            .unwrap_err();
708        assert_matches!(
709            err,
710            Error::TypeMismatch { expected, actual }
711                if expected == String::typename() && actual == TypedPayload::typehash().to_string()
712        );
713    }
714
715    #[test]
716    fn test_part_codec_uses_multipart_parts() {
717        let value = CodecPayload {
718            label: "codec".to_string(),
719            value: 123,
720        };
721        let repr = value.to_repr().unwrap();
722
723        let bincode_serialized = bincode::serialize(&value).unwrap();
724        let bincode_deserialized: CodecPayload = bincode::deserialize(&bincode_serialized).unwrap();
725        assert_eq!(bincode_deserialized, value);
726
727        let message = serialize_bincode(&value).unwrap();
728        assert!(message.body().is_empty());
729        assert_eq!(message.num_parts(), 1);
730        assert_eq!(
731            message.parts()[0].typehash(),
732            Some(CodecPayloadRepr::typehash())
733        );
734        assert_eq!(
735            message.parts()[0]
736                .deserialized::<CodecPayloadRepr>()
737                .unwrap(),
738            repr
739        );
740
741        let deserialized: CodecPayload = deserialize_bincode(message).unwrap();
742        assert_eq!(deserialized, value);
743    }
744
745    #[test]
746    fn test_part_codec_compound_multipart() {
747        #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
748        struct Envelope {
749            none: Option<CodecPayload>,
750            one: Option<CodecPayload>,
751            raw: Part,
752            many: Vec<CodecPayload>,
753            tail: u64,
754        }
755
756        let value = Envelope {
757            none: None,
758            one: Some(CodecPayload {
759                label: "one".to_string(),
760                value: 1,
761            }),
762            raw: Part::from("raw"),
763            many: vec![
764                CodecPayload {
765                    label: "two".to_string(),
766                    value: 2,
767                },
768                CodecPayload {
769                    label: "three".to_string(),
770                    value: 3,
771                },
772            ],
773            tail: 4,
774        };
775
776        let message = serialize_bincode(&value).unwrap();
777        assert_eq!(message.num_parts(), 4);
778        assert_eq!(
779            message.parts()[0].typehash(),
780            Some(CodecPayloadRepr::typehash())
781        );
782        assert_eq!(message.parts()[1].typehash(), None);
783        assert_eq!(message.parts()[1].to_bytes(), value.raw.to_bytes());
784        assert!(
785            message.parts()[2..]
786                .iter()
787                .all(|part| part.typehash() == Some(CodecPayloadRepr::typehash()))
788        );
789
790        let deserialized: Envelope = deserialize_bincode(message).unwrap();
791        assert_eq!(deserialized, value);
792    }
793
794    #[test]
795    fn test_visit_parts_mut() {
796        #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
797        struct Envelope {
798            first: CodecPayload,
799            raw_typed: Part,
800            raw_untyped: Part,
801            many: Vec<CodecPayload>,
802        }
803
804        let typed_payload = TypedPayload {
805            label: "typed".to_string(),
806            value: 7,
807        };
808        let value = Envelope {
809            first: CodecPayload {
810                label: "one".to_string(),
811                value: 1,
812            },
813            raw_typed: Part::serialize(&typed_payload).unwrap(),
814            raw_untyped: Part::from("raw"),
815            many: vec![
816                CodecPayload {
817                    label: "two".to_string(),
818                    value: 2,
819                },
820                CodecPayload {
821                    label: "three".to_string(),
822                    value: 3,
823                },
824            ],
825        };
826
827        let mut message = serialize_bincode(&value).unwrap();
828        assert_eq!(message.num_parts(), 5);
829
830        let mut visited = 0;
831        message
832            .visit_parts_mut::<CodecPayloadRepr, _>(|repr| {
833                visited += 1;
834                repr.label.push_str("-visited");
835                repr.value += 10;
836                Ok::<(), Error>(())
837            })
838            .unwrap();
839
840        assert_eq!(visited, 3);
841        assert_eq!(
842            message
843                .parts()
844                .iter()
845                .filter(|part| part.typehash() == Some(CodecPayloadRepr::typehash()))
846                .count(),
847            3
848        );
849        assert_eq!(
850            message.parts()[1].typehash(),
851            Some(TypedPayload::typehash())
852        );
853        assert_eq!(
854            message.parts()[1].deserialized::<TypedPayload>().unwrap(),
855            typed_payload
856        );
857        assert_eq!(message.parts()[2].typehash(), None);
858        assert_eq!(message.parts()[2].to_bytes(), value.raw_untyped.to_bytes());
859
860        let deserialized: Envelope = deserialize_bincode(message).unwrap();
861        assert_eq!(
862            deserialized,
863            Envelope {
864                first: CodecPayload {
865                    label: "one-visited".to_string(),
866                    value: 11,
867                },
868                raw_typed: value.raw_typed,
869                raw_untyped: value.raw_untyped,
870                many: vec![
871                    CodecPayload {
872                        label: "two-visited".to_string(),
873                        value: 12,
874                    },
875                    CodecPayload {
876                        label: "three-visited".to_string(),
877                        value: 13,
878                    },
879                ],
880            }
881        );
882    }
883
884    #[test]
885    fn test_part_codec_malformed_messages() {
886        let value = CodecPayload {
887            label: "codec".to_string(),
888            value: 123,
889        };
890        let mut message = serialize_bincode(&value).unwrap();
891        message.parts[0] = Part::serialize(&TypedPayload {
892            label: "wrong".to_string(),
893            value: 999,
894        })
895        .unwrap();
896        let err = deserialize_bincode::<CodecPayload>(message).unwrap_err();
897        assert_matches!(*err, bincode::ErrorKind::Custom(message) if message.contains("type mismatch"));
898
899        let mut message = serialize_bincode(&value).unwrap();
900        message.parts.clear();
901        let err = deserialize_bincode::<CodecPayload>(message).unwrap_err();
902        assert_matches!(*err, bincode::ErrorKind::Custom(message) if message == "multipart underrun while decoding");
903
904        let mut message = serialize_bincode(&42u64).unwrap();
905        message.parts.push(Part::from("extra"));
906        let err = deserialize_bincode::<u64>(message).unwrap_err();
907        assert_matches!(*err, bincode::ErrorKind::Custom(message) if message == "multipart overrun while decoding");
908    }
909
910    #[test]
911    fn test_part_codec_macro() {
912        let value = MacroCodecPayload {
913            label: "macro".to_string(),
914            value: 456,
915        };
916
917        let bincode_serialized = bincode::serialize(&value).unwrap();
918        let bincode_deserialized: MacroCodecPayload =
919            bincode::deserialize(&bincode_serialized).unwrap();
920        assert_eq!(bincode_deserialized, value);
921
922        let message = serialize_bincode(&value).unwrap();
923        assert!(message.body().is_empty());
924        assert_eq!(message.num_parts(), 1);
925        assert_eq!(
926            message.parts()[0].typehash(),
927            Some(MacroCodecPayloadRepr::typehash())
928        );
929
930        let deserialized: MacroCodecPayload = deserialize_bincode(message).unwrap();
931        assert_eq!(deserialized, value);
932    }
933
934    #[test]
935    fn test_raw_part_envelope_deserializes_payload() {
936        #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
937        struct Envelope {
938            header: String,
939            payload: Part,
940        }
941
942        let payload = TypedPayload {
943            label: "payload".to_string(),
944            value: 7,
945        };
946        let envelope = Envelope {
947            header: "header".to_string(),
948            payload: Part::serialize(&payload).unwrap(),
949        };
950
951        let message = serialize_bincode(&envelope).unwrap();
952        assert_eq!(message.num_parts(), 1);
953
954        let deserialized: Envelope = deserialize_bincode(message).unwrap();
955        assert_eq!(deserialized.header, envelope.header);
956        assert_eq!(
957            deserialized.payload.deserialized::<TypedPayload>().unwrap(),
958            payload
959        );
960    }
961
962    #[test]
963    fn test_malformed_messages() {
964        let message = Message {
965            body: Part::from("hello"),
966            parts: vec![Part::from("world")],
967        };
968        let err = deserialize_bincode::<String>(message).unwrap_err();
969
970        // Normal bincode errors work:
971        assert_matches!(*err, bincode::ErrorKind::Io(err) if err.kind() == std::io::ErrorKind::UnexpectedEof);
972
973        let mut message =
974            serialize_bincode(&vec![Part::from("hello"), Part::from("world")]).unwrap();
975        message.parts.push(Part::from("foo"));
976        let err = deserialize_bincode::<Vec<Part>>(message).unwrap_err();
977        assert_matches!(*err, bincode::ErrorKind::Custom(message) if message == "multipart overrun while decoding");
978
979        let mut message =
980            serialize_bincode(&vec![Part::from("hello"), Part::from("world")]).unwrap();
981        let _dropped_message = message.parts.pop().unwrap();
982        let err = deserialize_bincode::<Vec<Part>>(message).unwrap_err();
983        assert_matches!(*err, bincode::ErrorKind::Custom(message) if message == "multipart underrun while decoding");
984    }
985
986    #[test]
987    fn test_concat_buf() {
988        let buffers = vec![
989            Bytes::from("hello"),
990            Bytes::from("world"),
991            Bytes::from("1"),
992            Bytes::from(""),
993            Bytes::from("xyz"),
994            Bytes::from("xyzd"),
995        ];
996
997        let mut concat = Frame::from_buffers(buffers.clone());
998
999        assert_eq!(concat.remaining(), 18);
1000        concat.advance(2);
1001        assert_eq!(concat.remaining(), 16);
1002        assert_eq!(concat.chunk(), &b"llo"[..]);
1003        concat.advance(4);
1004        assert_eq!(concat.chunk(), &b"orld"[..]);
1005        concat.advance(5);
1006        assert_eq!(concat.chunk(), &b"xyz"[..]);
1007
1008        let mut concat = Frame::from_buffers(buffers);
1009        let bytes = concat.copy_to_bytes(concat.remaining());
1010        assert_eq!(&*bytes, &b"helloworld1xyzxyzd"[..]);
1011    }
1012
1013    #[test]
1014    fn test_framing() {
1015        let message = Message {
1016            body: Part::from("hello"),
1017            parts: vec![
1018                Part::from("world"),
1019                Part::from("1"),
1020                Part::from(""),
1021                Part::from("xyz"),
1022                Part::from("xyzd"),
1023            ],
1024        };
1025
1026        let mut framed = message.clone().framed();
1027        let framed = framed.copy_to_bytes(framed.remaining());
1028        assert_eq!(Message::from_framed(framed).unwrap(), message);
1029    }
1030
1031    #[test]
1032    fn test_typed_framing() {
1033        let body_value = TypedPayload {
1034            label: "body".to_string(),
1035            value: 1,
1036        };
1037        let part_value = TypedPayload {
1038            label: "part".to_string(),
1039            value: 2,
1040        };
1041        let body = Part::serialize(&body_value).unwrap();
1042        let typed_part = Part::serialize(&part_value).unwrap();
1043        let opaque_part = Part::from("opaque");
1044        let message = Message::from_body_and_parts(
1045            body.clone(),
1046            vec![opaque_part.clone(), typed_part.clone()],
1047        );
1048
1049        let mut frame = message.clone().framed();
1050        assert_eq!(frame.remaining(), message.frame_len());
1051        let encoded = frame.copy_to_bytes(frame.remaining());
1052
1053        let mut bytes = encoded.clone();
1054        assert_eq!(bytes.get_u64(), FRAME_TYPEHASH_FLAG | body.len() as u64);
1055        assert_eq!(bytes.get_u64(), TypedPayload::typehash());
1056        assert_eq!(bytes.split_to(body.len()), body.to_bytes());
1057        assert_eq!(bytes.get_u64(), opaque_part.len() as u64);
1058        assert_eq!(bytes.split_to(opaque_part.len()), opaque_part.to_bytes());
1059        assert_eq!(
1060            bytes.get_u64(),
1061            FRAME_TYPEHASH_FLAG | typed_part.len() as u64
1062        );
1063        assert_eq!(bytes.get_u64(), TypedPayload::typehash());
1064        assert_eq!(bytes.split_to(typed_part.len()), typed_part.to_bytes());
1065        assert!(bytes.is_empty());
1066
1067        let unframed = Message::from_framed(encoded).unwrap();
1068        assert_eq!(unframed, message);
1069        assert_eq!(
1070            unframed.body().deserialized::<TypedPayload>().unwrap(),
1071            body_value
1072        );
1073        assert_eq!(
1074            unframed.parts()[1].deserialized::<TypedPayload>().unwrap(),
1075            part_value
1076        );
1077    }
1078
1079    #[test]
1080    fn test_typed_framing_malformed_messages() {
1081        let mut frame = BytesMut::new();
1082        frame.put_u64(FRAME_TYPEHASH_FLAG | 1);
1083        let err = Message::from_framed(frame.freeze()).unwrap_err();
1084        assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
1085
1086        let mut frame = BytesMut::new();
1087        frame.put_u64(FRAME_TYPEHASH_FLAG | 4);
1088        frame.put_u64(TypedPayload::typehash());
1089        frame.extend_from_slice(b"xx");
1090        let err = Message::from_framed(frame.freeze()).unwrap_err();
1091        assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof);
1092    }
1093
1094    #[test]
1095    fn test_framing_rejects_lengths_that_need_typehash_bit() {
1096        let Some(too_large) = usize::try_from(FRAME_LEN_MASK)
1097            .ok()
1098            .and_then(|len| len.checked_add(1))
1099        else {
1100            return;
1101        };
1102
1103        let result = std::panic::catch_unwind(|| Message::frame_tag(too_large, None));
1104        assert!(result.is_err());
1105    }
1106
1107    #[test]
1108    fn test_socket_addr() {
1109        let socket_addr_v6: SocketAddrV6 =
1110            "[2401:db00:225c:2d09:face:0:223:0]:48483".parse().unwrap();
1111        {
1112            let message = serialize_bincode(&socket_addr_v6).unwrap();
1113            let deserialized: SocketAddrV6 = deserialize_bincode(message).unwrap();
1114            assert_eq!(socket_addr_v6, deserialized);
1115        }
1116        let socket_addr = SocketAddr::V6(socket_addr_v6);
1117        {
1118            let message = serialize_bincode(&socket_addr).unwrap();
1119            let deserialized: SocketAddr = deserialize_bincode(message).unwrap();
1120            assert_eq!(socket_addr, deserialized);
1121        }
1122
1123        let mut address_book: HashMap<usize, SocketAddr> = HashMap::new();
1124        address_book.insert(1, socket_addr);
1125        {
1126            let message = serialize_bincode(&address_book).unwrap();
1127            let deserialized: HashMap<usize, SocketAddr> = deserialize_bincode(message).unwrap();
1128            assert_eq!(address_book, deserialized);
1129        }
1130    }
1131
1132    prop_compose! {
1133        fn arb_bytes()(len in 0..1000000usize) -> Bytes {
1134            Bytes::from(vec![42; len])
1135        }
1136    }
1137
1138    prop_compose! {
1139        fn arb_part()(bytes in arb_bytes()) -> Part {
1140            bytes.into()
1141        }
1142    }
1143
1144    #[derive(Arbitrary, Serialize, Deserialize, Debug, PartialEq)]
1145    enum TupleEnum {
1146        One,
1147        Two(String),
1148        Three(u32),
1149    }
1150
1151    #[derive(Arbitrary, Serialize, Deserialize, Debug, PartialEq)]
1152    enum StructEnum {
1153        One {
1154            a: i32,
1155        },
1156        Two {
1157            s: String,
1158        },
1159        Three {
1160            e: TupleEnum,
1161            s: String,
1162            u: u32,
1163        },
1164        Four {
1165            #[proptest(strategy = "arb_part()")]
1166            part: Part,
1167        },
1168    }
1169
1170    #[derive(Arbitrary, Serialize, Deserialize, Debug, PartialEq)]
1171    struct S {
1172        field: String,
1173        tup: (StructEnum, i32, String, u32, f32),
1174        tup2: Option<(String, String, String, i32)>,
1175        e: StructEnum,
1176        maybe_e: Option<StructEnum>,
1177        many_e: Vec<(StructEnum, Option<TupleEnum>)>,
1178        #[proptest(strategy = "arb_bytes()")]
1179        some_bytes: Bytes,
1180    }
1181
1182    #[derive(Arbitrary, Serialize, Deserialize, Debug, PartialEq)]
1183    struct N(S);
1184
1185    proptest! {
1186        #[test]
1187        fn test_arbitrary_roundtrip(value in any::<N>()) {
1188            let message = serialize_bincode(&value).unwrap();
1189            let deserialized_value = deserialize_bincode(message.clone()).unwrap();
1190            assert_eq!(value, deserialized_value);
1191        }
1192    }
1193}