ergot_base/interface_manager/
std_utils.rs

1use crate::{Address, FrameKind, HeaderSeq, Key, ProtocolError};
2
3pub(crate) struct OwnedFrame {
4    pub(crate) hdr: HeaderSeq,
5    pub(crate) body: Result<Vec<u8>, ProtocolError>,
6}
7
8#[derive(Debug, PartialEq)]
9pub enum ReceiverError {
10    SocketClosed,
11}
12
13pub(crate) fn ser_frame(frame: OwnedFrame) -> Vec<u8> {
14    let dst_any = [0, 255].contains(&frame.hdr.dst.port_id);
15    if frame.hdr.kind == FrameKind::PROTOCOL_ERROR {
16        assert!(!dst_any);
17        assert!(frame.body.is_err());
18    }
19    let src = frame.hdr.src.as_u32();
20    let dst = frame.hdr.dst.as_u32();
21    let seq = frame.hdr.seq_no;
22
23    let mut out = vec![];
24    // TODO: This is bad and does a ton of allocs. yolo
25    //
26    out.extend_from_slice(&postcard::to_stdvec(&src).unwrap());
27    out.extend_from_slice(&postcard::to_stdvec(&dst).unwrap());
28    out.push(frame.hdr.kind.0);
29    out.push(frame.hdr.ttl);
30    if dst_any {
31        let key = frame.hdr.key.unwrap();
32        out.extend_from_slice(&key.0);
33    }
34
35    out.extend_from_slice(&postcard::to_stdvec(&seq).unwrap());
36    match frame.body {
37        Ok(body) => {
38            assert!(frame.hdr.kind != FrameKind::PROTOCOL_ERROR);
39            out.extend_from_slice(&body)
40        }
41        Err(err) => {
42            assert!(frame.hdr.kind == FrameKind::PROTOCOL_ERROR);
43            out.extend_from_slice(&postcard::to_stdvec(&err).unwrap())
44        }
45    }
46    let mut out = cobs::encode_vec(&out);
47    out.push(0);
48    out
49}
50
51pub(crate) fn de_frame(remain: &[u8]) -> Option<OwnedFrame> {
52    let (src_word, remain) = postcard::take_from_bytes::<u32>(remain).ok()?;
53    let src = Address::from_word(src_word);
54    let (dst_word, remain) = postcard::take_from_bytes::<u32>(remain).ok()?;
55    let dst = Address::from_word(dst_word);
56    let (kind, remain) = remain.split_first()?;
57    let kind = FrameKind(*kind);
58    let (ttl, remain) = remain.split_first()?;
59    let ttl = *ttl;
60
61    let is_err = kind == FrameKind::PROTOCOL_ERROR;
62
63    let (key, remain) = if [0, 255].contains(&dst.port_id) {
64        assert!(!is_err);
65        if remain.len() < 8 {
66            return None;
67        }
68        let (keyb, remain) = remain.split_at(8);
69        let mut buf = [0u8; 8];
70        buf.copy_from_slice(keyb);
71        (Some(Key(buf)), remain)
72    } else {
73        (None, remain)
74    };
75
76    let (seq, remain) = postcard::take_from_bytes::<u16>(remain).ok()?;
77
78    let body = if is_err {
79        let (err, remain) = postcard::take_from_bytes::<ProtocolError>(remain).ok()?;
80        assert!(remain.is_empty());
81        Err(err)
82    } else {
83        Ok(remain.to_vec())
84    };
85
86    Some(OwnedFrame {
87        hdr: HeaderSeq {
88            src,
89            dst,
90            seq_no: seq,
91            key,
92            kind,
93            ttl,
94        },
95        body,
96    })
97}
98
99pub(crate) mod acc {
100    //! Basically postcard's cobs accumulator, but without the deser part
101
102    pub struct CobsAccumulator {
103        buf: Box<[u8]>,
104        idx: usize,
105    }
106
107    /// The result of feeding the accumulator.
108    pub enum FeedResult<'input, 'buf> {
109        /// Consumed all data, still pending.
110        Consumed,
111
112        /// Buffer was filled. Contains remaining section of input, if any.
113        OverFull(&'input [u8]),
114
115        /// Reached end of chunk, but deserialization failed. Contains remaining section of input, if.
116        /// any
117        DeserError(&'input [u8]),
118
119        Success {
120            /// Decoded data.
121            data: &'buf [u8],
122
123            /// Remaining data left in the buffer after deserializing.
124            remaining: &'input [u8],
125        },
126    }
127
128    impl CobsAccumulator {
129        /// Create a new accumulator.
130        pub fn new(sz: usize) -> Self {
131            CobsAccumulator {
132                buf: vec![0u8; sz].into_boxed_slice(),
133                idx: 0,
134            }
135        }
136
137        /// Appends data to the internal buffer and attempts to deserialize the accumulated data into
138        /// `T`.
139        ///
140        /// This differs from feed, as it allows the `T` to reference data within the internal buffer, but
141        /// mutably borrows the accumulator for the lifetime of the deserialization.
142        /// If `T` does not require the reference, the borrow of `self` ends at the end of the function.
143        pub fn feed_raw<'me, 'input>(
144            &'me mut self,
145            input: &'input [u8],
146        ) -> FeedResult<'input, 'me> {
147            if input.is_empty() {
148                return FeedResult::Consumed;
149            }
150
151            let zero_pos = input.iter().position(|&i| i == 0);
152            let max_len = self.buf.len();
153
154            if let Some(n) = zero_pos {
155                // Yes! We have an end of message here.
156                // Add one to include the zero in the "take" portion
157                // of the buffer, rather than in "release".
158                let (take, release) = input.split_at(n + 1);
159
160                // TODO(AJM): We could special case when idx == 0 to avoid copying
161                // into the dest buffer if there's a whole packet in the input
162
163                // Does it fit?
164                if (self.idx + take.len()) <= max_len {
165                    // Aw yiss - add to array
166                    self.extend_unchecked(take);
167
168                    let retval = match cobs::decode_in_place(&mut self.buf[..self.idx]) {
169                        Ok(ct) => FeedResult::Success {
170                            data: &self.buf[..ct],
171                            remaining: release,
172                        },
173                        Err(_) => FeedResult::DeserError(release),
174                    };
175                    self.idx = 0;
176                    retval
177                } else {
178                    self.idx = 0;
179                    FeedResult::OverFull(release)
180                }
181            } else {
182                // Does it fit?
183                if (self.idx + input.len()) > max_len {
184                    // nope
185                    let new_start = max_len - self.idx;
186                    self.idx = 0;
187                    FeedResult::OverFull(&input[new_start..])
188                } else {
189                    // yup!
190                    self.extend_unchecked(input);
191                    FeedResult::Consumed
192                }
193            }
194        }
195
196        /// Extend the internal buffer with the given input.
197        ///
198        /// # Panics
199        ///
200        /// Will panic if the input does not fit in the internal buffer.
201        fn extend_unchecked(&mut self, input: &[u8]) {
202            let new_end = self.idx + input.len();
203            self.buf[self.idx..new_end].copy_from_slice(input);
204            self.idx = new_end;
205        }
206    }
207}