Skip to main content

cliquenet/msg/
trailer.rs

1use bytes::Bytes;
2
3use crate::msg::{MsgId, Slot};
4
5const STD: u8 = 0;
6const NO_ACK: u8 = 1;
7
8/// Meta information appended at the end of a message.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum Trailer {
11    Std {
12        /// The slot the message corresponds to.
13        slot: Slot,
14        /// The message ID.
15        id: MsgId,
16    },
17    NoAck {
18        /// The slot the message corresponds to.
19        slot: Slot,
20    },
21    Unknown,
22}
23
24impl Trailer {
25    pub const MAX_SIZE: usize = u8::MAX as usize;
26
27    pub(crate) fn from_bytes(bytes: &mut Bytes) -> Option<Self> {
28        if bytes.len() < 2 {
29            return None;
30        }
31        let len = bytes.len();
32        let trailer_len = bytes[len - 1];
33        let trailer_typ = bytes[len - 2];
34        match trailer_typ {
35            STD => {
36                if trailer_len != 16 || len < 18 {
37                    return None;
38                }
39                let id = u64::from_be_bytes(bytes[len - 10..len - 2].try_into().ok()?);
40                let slot = u64::from_be_bytes(bytes[len - 18..len - 10].try_into().ok()?);
41                bytes.truncate(len - 18);
42                Some(Self::Std {
43                    slot: Slot(slot),
44                    id: MsgId(id),
45                })
46            },
47            NO_ACK => {
48                if trailer_len != 8 || len < 10 {
49                    return None;
50                }
51                let slot = u64::from_be_bytes(bytes[len - 10..len - 2].try_into().ok()?);
52                bytes.truncate(len - 10);
53                Some(Self::NoAck { slot: Slot(slot) })
54            },
55            _ => {
56                let trailer_len = 2 + usize::from(trailer_len);
57                if trailer_len > len {
58                    return None;
59                }
60                bytes.truncate(len - trailer_len);
61                Some(Self::Unknown)
62            },
63        }
64    }
65
66    pub(crate) fn to_bytes(self) -> TrailerBytes {
67        match self {
68            Self::Std { slot, id } => {
69                let mut buf = [0; 18];
70                buf[..8].copy_from_slice(&slot.0.to_be_bytes()[..]);
71                buf[8..16].copy_from_slice(&id.0.to_be_bytes()[..]);
72                buf[16] = STD;
73                buf[17] = 16;
74                TrailerBytes::Std(buf)
75            },
76            Self::NoAck { slot } => {
77                let mut buf = [0; 10];
78                buf[..8].copy_from_slice(&slot.0.to_be_bytes()[..]);
79                buf[8] = NO_ACK;
80                buf[9] = 8;
81                TrailerBytes::NoAck(buf)
82            },
83            Self::Unknown => {
84                unreachable!("nothing constructs an unknown trailer")
85            },
86        }
87    }
88}
89
90#[derive(Debug, Clone, Copy)]
91pub(crate) enum TrailerBytes {
92    Std([u8; 18]),
93    NoAck([u8; 10]),
94}
95
96impl AsRef<[u8]> for TrailerBytes {
97    fn as_ref(&self) -> &[u8] {
98        match self {
99            Self::Std(a) => &a[..],
100            Self::NoAck(a) => &a[..],
101        }
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use bytes::Bytes;
108    use quickcheck::{Arbitrary, Gen, quickcheck};
109
110    use super::{NO_ACK, STD, Trailer};
111    use crate::msg::{MsgId, Slot};
112
113    impl Arbitrary for Trailer {
114        fn arbitrary(g: &mut Gen) -> Self {
115            match bool::arbitrary(g) {
116                true => Self::Std {
117                    slot: Slot(u64::arbitrary(g)),
118                    id: MsgId(u64::arbitrary(g)),
119                },
120                false => Self::NoAck {
121                    slot: Slot(u64::arbitrary(g)),
122                },
123            }
124        }
125    }
126
127    /// Trailer-shaped bytes: a short body plus a tail biased toward the type
128    /// and length bytes the parser recognises, so the short-input slicing
129    /// paths are hit on nearly every case (uniform `Vec<u8>` almost never is).
130    #[derive(Debug, Clone)]
131    struct TrailerLike(Vec<u8>);
132
133    impl Arbitrary for TrailerLike {
134        fn arbitrary(g: &mut Gen) -> Self {
135            let body_len = usize::arbitrary(g) % 20;
136            let mut bytes: Vec<u8> = (0..body_len).map(|_| u8::arbitrary(g)).collect();
137
138            let rand_typ = u8::arbitrary(g);
139            let typ = *g.choose(&[STD, NO_ACK, rand_typ]).unwrap();
140            let rand_len = u8::arbitrary(g);
141            let trailer_len = *g.choose(&[16u8, 8, rand_len]).unwrap();
142            bytes.push(typ);
143            bytes.push(trailer_len);
144            TrailerLike(bytes)
145        }
146    }
147
148    quickcheck! {
149        fn prop_to_bytes_from_bytes_id(t1: Trailer) -> bool {
150            let mut b = Bytes::copy_from_slice(t1.to_bytes().as_ref());
151            let t2 = Trailer::from_bytes(&mut b);
152            Some(t1) == t2
153        }
154
155        fn prop_from_bytes_arbitrary_does_not_panic(bytes: Vec<u8>) -> bool {
156            let mut b = Bytes::copy_from_slice(&bytes);
157            let _ = Trailer::from_bytes(&mut b);
158            true
159        }
160
161        fn prop_from_bytes_trailer_like_does_not_panic(t: TrailerLike) -> bool {
162            let mut b = Bytes::copy_from_slice(&t.0);
163            let _ = Trailer::from_bytes(&mut b);
164            true
165        }
166    }
167
168    #[test]
169    fn regression_std_trailer_short_input_does_not_panic() {
170        let mut b = Bytes::copy_from_slice(&[STD, 16]);
171        assert_eq!(Trailer::from_bytes(&mut b), None);
172    }
173
174    #[test]
175    fn regression_no_ack_trailer_short_input_does_not_panic() {
176        let mut b = Bytes::copy_from_slice(&[NO_ACK, 8]);
177        assert_eq!(Trailer::from_bytes(&mut b), None);
178    }
179}