asynchronous_codec/codec/
length.rs

1use crate::{Decoder, Encoder};
2use bytes::{Buf, BufMut, Bytes, BytesMut};
3use std::io::Error;
4
5const U64_LENGTH: usize = std::mem::size_of::<u64>();
6
7/// A simple `Codec` implementation sending your data by prefixing it by its length.
8///
9/// # Example
10///
11/// This codec will most likely be used wrapped in another codec like so.
12///
13/// ```
14/// use asynchronous_codec::{Decoder, Encoder, LengthCodec};
15/// use bytes::{Bytes, BytesMut};
16/// use std::io::{Error, ErrorKind};
17///
18/// pub struct MyStringCodec(LengthCodec);
19///
20/// impl Encoder for MyStringCodec {
21///     type Item = String;
22///     type Error = Error;
23///
24///     fn encode(&mut self, src: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
25///         let bytes = Bytes::from(src);
26///         self.0.encode(bytes, dst)
27///     }
28/// }
29///
30/// impl Decoder for MyStringCodec {
31///     type Item = String;
32///     type Error = Error;
33///
34///     fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
35///         match self.0.decode(src)? {
36///             Some(bytes) => {
37///                 match String::from_utf8(bytes.to_vec()) {
38///                     Ok(string) => Ok(Some(string)),
39///                     Err(e) => Err(Error::new(ErrorKind::InvalidData, e))
40///                 }
41///             },
42///             None => Ok(None),
43///         }
44///     }
45/// }
46/// ```
47pub struct LengthCodec;
48
49impl Encoder for LengthCodec {
50    type Item = Bytes;
51    type Error = Error;
52
53    fn encode(&mut self, src: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> {
54        dst.reserve(U64_LENGTH + src.len());
55        dst.put_u64(src.len() as u64);
56        dst.extend_from_slice(&src);
57        Ok(())
58    }
59}
60
61impl Decoder for LengthCodec {
62    type Item = Bytes;
63    type Error = Error;
64
65    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
66        if src.len() < U64_LENGTH {
67            return Ok(None);
68        }
69
70        let mut len_bytes = [0u8; U64_LENGTH];
71        len_bytes.copy_from_slice(&src[..U64_LENGTH]);
72        let len = u64::from_be_bytes(len_bytes) as usize;
73
74        if src.len() - U64_LENGTH >= len {
75            // Skip the length header we already read.
76            src.advance(U64_LENGTH);
77            Ok(Some(src.split_to(len).freeze()))
78        } else {
79            Ok(None)
80        }
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    mod decode {
89        use super::*;
90
91        #[test]
92        fn it_returns_bytes_withouth_length_header() {
93            let mut codec = LengthCodec {};
94
95            let mut src = BytesMut::with_capacity(5);
96            src.put(&[0, 0, 0, 0, 0, 0, 0, 3u8, 1, 2, 3, 4][..]);
97            let item = codec.decode(&mut src).unwrap();
98
99            assert!(item == Some(Bytes::from(&[1u8, 2, 3][..])));
100        }
101    }
102}