diff --git a/src/proto/base_codec.rs b/src/proto/base_codec.rs index eeed294..1c57755 100644 --- a/src/proto/base_codec.rs +++ b/src/proto/base_codec.rs @@ -11,7 +11,7 @@ use encoding::{DecoderTrap, EncoderTrap, Encoding}; // --------- /// Length of an encoded 32-bit integer in bytes. -const U32_BYTE_LEN: usize = 4; +pub const U32_BYTE_LEN: usize = 4; /*===================================* * BASIC TYPES ENCODING AND DECODING * diff --git a/src/proto/codec.rs b/src/proto/codec.rs index 7863430..bb5a8d7 100644 --- a/src/proto/codec.rs +++ b/src/proto/codec.rs @@ -4,7 +4,7 @@ use std::marker; use tokio_codec; use bytes::BytesMut; -use super::base_codec::{Decode, ProtoEncode, ProtoEncoder}; +use super::base_codec::{Decode, ProtoEncode, ProtoEncoder, U32_BYTE_LEN}; use super::server::{ServerRequest,ServerResponse}; use super::peer::Message; @@ -12,8 +12,13 @@ use super::peer::Message; * TOKIO CODEC TRAIT IMPLEMENTATIONS * *===================================*/ +// Encodes types that implement ProtoEncode with a length prefix. struct Encoder { - data: marker::PhantomData + phantom: marker::PhantomData +} + +impl Encoder { + fn new() -> Self { Self{phantom: marker::PhantomData} } } impl tokio_codec::Encoder for Encoder { @@ -21,8 +26,18 @@ impl tokio_codec::Encoder for Encoder { type Error = io::Error; fn encode(&mut self, item: Self::Item, dst: &mut BytesMut) -> Result<(), Self::Error> { - let mut encoder = ProtoEncoder::new(dst); - item.encode(&mut encoder)?; + // Split buffer into two parts: the length prefix and the message. + dst.reserve(U32_BYTE_LEN); + let mut msg_dst = dst.split_off(U32_BYTE_LEN); + + // Encode the message. + item.encode(&mut ProtoEncoder::new(&mut msg_dst))?; + + // Encode the message length. + ProtoEncoder::new(dst).encode_u32(msg_dst.len() as u32)?; + + // Reassemble both parts into one contiguous buffer. + dst.unsplit(msg_dst); Ok(()) } } @@ -48,3 +63,43 @@ where BytesMut: Decode { pub type ServerRequestDecoder = Decoder; pub type ServerResponseDecoder = Decoder; pub type PeerMessageDecoder = Decoder; + +mod tests { + use bytes::BytesMut; + use tokio_codec::Encoder; + + use proto::ProtoEncode; + + // Avoid name conflict with tokio_codec::Encoder. + use super::Encoder as MyEncoder; + + #[test] + fn encode_u32() { + let val: u32 = 13 + 37*256; + + let mut bytes = BytesMut::new(); + MyEncoder::new().encode(val, &mut bytes).unwrap(); + + assert_eq!(bytes, vec![ + 4, 0, 0, 0, + 13, 37, 0, 0, + ]); + } + + #[test] + fn encode_vec() { + let v: Vec = vec![1, 3, 3, 7]; + + let mut bytes = BytesMut::new(); + MyEncoder::new().encode(v, &mut bytes).unwrap(); + + assert_eq!(bytes, vec![ + 20, 0, 0, 0, // 5 32-bit integers = 20 bytes. + 4, 0, 0, 0, + 1, 0, 0, 0, + 3, 0, 0, 0, + 3, 0, 0, 0, + 7, 0, 0, 0 + ]); + } +}