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}