Implement length-prefixed encoding for Tokio.

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 *

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<T> {
-  data: marker::PhantomData<T>
+  phantom: marker::PhantomData<T>
+}
+
+impl<T> Encoder<T> {
+    fn new() -> Self { Self{phantom: marker::PhantomData} }
 }
 
 impl<T: ProtoEncode> tokio_codec::Encoder for Encoder<T> {
@@ -21,8 +26,18 @@ impl<T: ProtoEncode> tokio_codec::Encoder for Encoder<T> {
     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<T> {
 pub type ServerRequestDecoder = Decoder<ServerRequest>;
 pub type ServerResponseDecoder = Decoder<ServerResponse>;
 pub type PeerMessageDecoder = Decoder<Message>;
+
+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<u32> = 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
+        ]);
+    }
+}