compio_io/framed/
frame.rs

1//! Traits and implementations for frame extraction and enclosing
2
3use std::io;
4
5use compio_buf::{IoBuf, IoBufMut, Slice};
6
7/// An extracted frame
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub struct Frame {
10    /// Offset where the frame payload begins
11    prefix: usize,
12
13    /// Length of the frame payload
14    payload: usize,
15
16    /// Suffix length of the frame
17    suffix: usize,
18}
19
20impl Frame {
21    /// Create a new [`Frame`] with the specified prefix, payload, and suffix
22    /// lengths.
23    pub fn new(prefix: usize, payload: usize, suffix: usize) -> Self {
24        Self {
25            prefix,
26            payload,
27            suffix,
28        }
29    }
30
31    /// Length of the entire frame
32    pub fn len(&self) -> usize {
33        self.prefix + self.payload + self.suffix
34    }
35
36    /// If the frame is empty
37    pub fn is_empty(&self) -> bool {
38        self.len() == 0
39    }
40
41    /// Slice payload out of the buffer
42    pub fn slice<B: IoBuf>(&self, buf: B) -> Slice<B> {
43        buf.slice(self.prefix..self.prefix + self.payload)
44    }
45}
46
47/// Enclosing and extracting frames in a buffer.
48pub trait Framer<B: IoBufMut> {
49    /// Enclose a frame in the given buffer.
50    ///
51    /// All initialized bytes in `buf` (`buf[0..buf.buf_len()]`) are valid and
52    /// required to be enclosed. All modifications should happen in-place; one
53    /// can use [`IoBufMut::reserve`], [`IoBufMut::copy_within`] or a temporary
54    /// buffer if prepending data is necessary.
55    ///
56    /// [`slice::copy_within`]: https://doc.rust-lang.org/std/primitive.slice.html#method.copy_within
57    fn enclose(&mut self, buf: &mut B);
58
59    /// Extract a frame from the given buffer.
60    ///
61    /// # Returns
62    /// - `Ok(Some(frame))` if a complete frame is found.
63    /// - `Ok(None)` if no complete frame is found.
64    /// - `Err(io::Error)` if an error occurs during extraction.
65    fn extract(&mut self, buf: &Slice<B>) -> io::Result<Option<Frame>>;
66}
67
68/// A simple extractor that frames data by its length.
69///
70/// It uses 8 bytes to represent the length of the data at the beginning.
71#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
72pub struct LengthDelimited {
73    length_field_len: usize,
74    length_field_is_big_endian: bool,
75}
76
77impl Default for LengthDelimited {
78    fn default() -> Self {
79        Self {
80            length_field_len: 4,
81            length_field_is_big_endian: true,
82        }
83    }
84}
85
86impl LengthDelimited {
87    /// Max allowed length of `len` field
88    const MAX_LFL: usize = 8;
89
90    /// Creates a new `LengthDelimited` framer.
91    pub fn new() -> Self {
92        Self::default()
93    }
94
95    /// Returns the length of the length field in bytes.
96    pub fn length_field_len(&self) -> usize {
97        self.length_field_len
98    }
99
100    /// Sets the length of the length field in bytes.
101    ///
102    /// # Panics
103    ///
104    /// This will panic if `len_field_len` is too long (8 bytes is the maximum
105    /// allowed for now).
106    pub fn set_length_field_len(mut self, len_field_len: usize) -> Self {
107        assert!(
108            len_field_len <= Self::MAX_LFL,
109            "Length field cannot take over 8 bytes"
110        );
111        self.length_field_len = len_field_len;
112        self
113    }
114
115    /// Returns whether the length field is big-endian.
116    pub fn length_field_is_big_endian(&self) -> bool {
117        self.length_field_is_big_endian
118    }
119
120    /// Sets whether the length field is big-endian.
121    pub fn set_length_field_is_big_endian(mut self, big_endian: bool) -> Self {
122        self.length_field_is_big_endian = big_endian;
123        self
124    }
125}
126
127impl<B: IoBufMut> Framer<B> for LengthDelimited {
128    fn enclose(&mut self, buf: &mut B) {
129        let len = (*buf).buf_len();
130
131        buf.reserve(self.length_field_len).expect("Reserve failed");
132        buf.copy_within(0..len, self.length_field_len); // Shift existing data
133        unsafe { buf.advance_to(len + self.length_field_len) };
134
135        let slice = buf.as_mut_slice();
136        let lfl = self.length_field_len;
137
138        // Write the length at the beginning
139        let len = len as u64;
140        let len_bytes = if self.length_field_is_big_endian {
141            &len.to_be_bytes()[Self::MAX_LFL - lfl..]
142        } else {
143            &len.to_le_bytes()[..lfl]
144        };
145        slice[..lfl].copy_from_slice(len_bytes);
146    }
147
148    fn extract(&mut self, buf: &Slice<B>) -> io::Result<Option<Frame>> {
149        if buf.len() < self.length_field_len {
150            return Ok(None);
151        }
152
153        let buf = buf.as_init();
154        let lfl = self.length_field_len;
155        let mut len_bytes = [0; Self::MAX_LFL];
156
157        let len = if self.length_field_is_big_endian {
158            len_bytes[Self::MAX_LFL - lfl..].copy_from_slice(&buf[..lfl]);
159            u64::from_be_bytes(len_bytes)
160        } else {
161            len_bytes[..lfl].copy_from_slice(&buf[..lfl]);
162            u64::from_le_bytes(len_bytes)
163        } as usize;
164
165        if buf.len() < self.length_field_len + len {
166            return Ok(None);
167        }
168
169        Ok(Some(Frame::new(self.length_field_len, len, 0)))
170    }
171}
172
173/// A generic delimiter that uses a single character encoded as UTF-8.
174///
175/// If you need to use a multi-byte delimiter or other encodings, consider using
176/// [`AnyDelimited`].
177#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
178pub struct CharDelimited<const C: char> {
179    char_buf: [u8; 4],
180}
181
182impl<const C: char> CharDelimited<C> {
183    /// Creates a new `CharDelimited`
184    pub fn new() -> Self {
185        Self { char_buf: [0; 4] }
186    }
187
188    fn as_any_delimited(&mut self) -> AnyDelimited<'_> {
189        let bytes = C.encode_utf8(&mut self.char_buf).as_bytes();
190
191        AnyDelimited::new(bytes)
192    }
193}
194
195impl<B: IoBufMut, const C: char> Framer<B> for CharDelimited<C> {
196    fn enclose(&mut self, buf: &mut B) {
197        self.as_any_delimited().enclose(buf);
198    }
199
200    fn extract(&mut self, buf: &Slice<B>) -> io::Result<Option<Frame>> {
201        self.as_any_delimited().extract(buf)
202    }
203}
204
205/// A generic delimiter that uses any sequence of bytes as a delimiter.
206#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
207pub struct AnyDelimited<'a> {
208    bytes: &'a [u8],
209}
210
211impl<'a> AnyDelimited<'a> {
212    /// Creates a new `AnyDelimited` with the specified delimiter bytes.
213    pub fn new(bytes: &'a [u8]) -> Self {
214        Self { bytes }
215    }
216}
217
218impl<B: IoBufMut> Framer<B> for AnyDelimited<'_> {
219    fn extract(&mut self, buf: &Slice<B>) -> io::Result<Option<Frame>> {
220        if buf.is_empty() {
221            return Ok(None);
222        }
223
224        // Search for the first occurrence of any byte in `self.bytes`
225        // TODO(George-Miao): Optimize with memchr if performance is a concern
226        if let Some(pos) = buf
227            .windows(self.bytes.len())
228            .position(|window| window == self.bytes)
229        {
230            Ok(Some(Frame::new(0, pos, self.bytes.len())))
231        } else {
232            Ok(None)
233        }
234    }
235
236    fn enclose(&mut self, buf: &mut B) {
237        buf.extend_from_slice(self.bytes)
238            .expect("Failed to append delimiter");
239    }
240}
241
242/// Delimiter that uses newline characters (`\n`) as delimiters.
243pub type LineDelimited = CharDelimited<'\n'>;
244
245/// A framer that does nothing.
246///
247/// It simply reserves space in the buffer without adding any framing.
248#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
249pub struct NoopFramer {
250    max_size: usize,
251}
252
253impl Default for NoopFramer {
254    fn default() -> Self {
255        Self { max_size: 4096 }
256    }
257}
258
259impl NoopFramer {
260    /// Creates a new `NoopFramer` framer.
261    pub fn new() -> Self {
262        Self::default()
263    }
264
265    /// Returns the size of the capacity.
266    pub fn max_size(&self) -> usize {
267        self.max_size
268    }
269}
270
271impl<B: IoBufMut> Framer<B> for NoopFramer {
272    fn enclose(&mut self, _: &mut B) {}
273
274    fn extract(&mut self, buf: &Slice<B>) -> io::Result<Option<Frame>> {
275        if buf.is_empty() {
276            return Ok(None);
277        }
278
279        let len = if buf.len() < self.max_size {
280            buf.len()
281        } else {
282            self.max_size
283        };
284
285        Ok(Some(Frame::new(0, len, 0)))
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use compio_buf::{IntoInner, IoBufMut};
292
293    use super::*;
294
295    #[test]
296    fn test_length_delimited() {
297        let mut framer = LengthDelimited::new();
298
299        let mut buf = Vec::from(b"hello");
300        framer.enclose(&mut buf);
301        assert_eq!(&buf.as_slice()[..9], b"\x00\x00\x00\x05hello");
302
303        let buf = buf.slice(..);
304        let frame = framer.extract(&buf).unwrap().unwrap();
305        let buf = buf.into_inner();
306        assert_eq!(frame, Frame::new(4, 5, 0));
307        let payload = frame.slice(buf);
308        assert_eq!(payload.as_init(), b"hello");
309    }
310
311    #[test]
312    fn test_noop_framer() {
313        let mut framer = NoopFramer::new();
314
315        let mut buf = Vec::from(b"hello");
316        framer.enclose(&mut buf);
317        assert_eq!(&buf.as_slice()[..5], b"hello");
318
319        let buf = buf.slice(..);
320        let frame = framer.extract(&buf).unwrap().unwrap();
321        let buf = buf.into_inner();
322        assert_eq!(frame, Frame::new(0, 5, 0));
323        let payload = frame.slice(buf);
324        assert_eq!(payload.as_init(), b"hello");
325    }
326
327    #[test]
328    fn test_char_delimited() {
329        let mut framer = CharDelimited::<'ℝ'>::new();
330
331        let mut buf = Vec::new();
332        IoBufMut::extend_from_slice(&mut buf, b"hello").unwrap();
333        framer.enclose(&mut buf);
334        assert_eq!(buf.as_slice(), "helloℝ".as_init());
335
336        let buf = buf.slice(..);
337        let frame = framer.extract(&buf).unwrap().unwrap();
338        assert_eq!(frame, Frame::new(0, 5, 3));
339        let payload = frame.slice(buf);
340        assert_eq!(payload.as_init(), b"hello");
341    }
342}
compio_io/framed/frame.rs

compio_io/framed/
frame.rs
