1 use crate::{Decoder, Encoder};
2 use bytes::{Buf, BufMut, Bytes, BytesMut};
3 use std::io::Error;
4 
5 const 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 /// ```
47 pub struct LengthCodec;
48 
49 impl Encoder for LengthCodec {
50     type Item = Bytes;
51     type Error = Error;
52 
encode(&mut self, src: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error>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 
61 impl Decoder for LengthCodec {
62     type Item = Bytes;
63     type Error = Error;
64 
decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error>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)]
85 mod tests {
86     use super::*;
87 
88     mod decode {
89         use super::*;
90 
91         #[test]
it_returns_bytes_withouth_length_header()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 }
103