1use bytes::Bytes;
2
3use crate::msg::{MsgId, Slot};
4
5const STD: u8 = 0;
6const NO_ACK: u8 = 1;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum Trailer {
11 Std {
12 slot: Slot,
14 id: MsgId,
16 },
17 NoAck {
18 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 #[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}