Extract prefix encoding logic into Prefixer.

src/proto/codec.rs

@@ -19,7 +19,8 @@ use tokio_codec;
 
 use super::peer::Message;
 use super::server::{ServerRequest, ServerResponse};
-use super::value_codec::{ValueDecode, ValueDecoder, ValueEncode, ValueEncoder, U32_BYTE_LEN};
+use super::u32::U32_BYTE_LEN;
+use super::value_codec::{ValueDecode, ValueDecoder, ValueEncode, ValueEncoder};
 
 /// Implements tokio's Encoder trait for types that implement ValueEncode.
 pub struct LengthPrefixedEncoder<T> {

src/proto/mod.rs

@@ -3,8 +3,10 @@ mod constants;
 mod handler;
 mod packet;
 pub mod peer;
+mod prefix;
 pub mod server;
 mod stream;
+pub mod u32;
 mod user;
 mod value_codec;
 

src/proto/prefix.rs

@@ -0,0 +1,111 @@
+//! Utility module for writing length-prefixed values, the length of which is
+//! unknown before the value is encoded.
+
+use std::convert::TryFrom;
+
+use crate::proto::u32::{encode_u32, U32_BYTE_LEN};
+
+/// Encodes a length prefix in a buffer.
+pub struct Prefixer {
+    /// The position of the length prefix in the target buffer.
+    position: usize,
+}
+
+impl Prefixer {
+    /// Reserves space for the length prefix at the end of the given buffer.
+    ///
+    /// Returns a prefixer for writing the length later.
+    pub fn new(buffer: &mut Vec<u8>) -> Prefixer {
+        // Remember where we were.
+        let result = Prefixer {
+            position: buffer.len(),
+        };
+        // Reserve enough bytes to write the prefix into later.
+        buffer.extend_from_slice(&[0; U32_BYTE_LEN]);
+        result
+    }
+
+    /// Writes the length prefix into the given buffer in the reserved space.
+    ///
+    /// The given buffer must be the same one passed to new(), and should not
+    /// have been truncated since then.
+    ///
+    /// Panics if the buffer is not large enough to store the prefix.
+    ///
+    /// Returns `Err(length)` if `length`, the length of the suffix, is too
+    /// large to store in the reserved space.
+    pub fn finalize(self, buffer: &mut Vec<u8>) -> Result<(), usize> {
+        // The position at which the value should have been encoded.
+        let value_position = self.position + U32_BYTE_LEN;
+        assert!(buffer.len() >= value_position);
+
+        // Calculate the value's length, check it is not too large.
+        let length = buffer.len() - value_position;
+        let length_u32 = u32::try_from(length).map_err(|_| length)?;
+
+        // Write the length prefix into the reserved space.
+        let length_bytes = encode_u32(length_u32);
+        buffer[self.position..value_position].copy_from_slice(&length_bytes);
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Prefixer;
+
+    use std::convert::TryInto;
+
+    use crate::proto::u32::{decode_u32, U32_BYTE_LEN};
+
+    #[test]
+    fn new_reserves_space() {
+        let mut buffer = vec![13];
+
+        Prefixer::new(&mut buffer);
+
+        assert_eq!(buffer.len(), U32_BYTE_LEN + 1);
+        assert_eq!(buffer[0], 13);
+    }
+
+    #[test]
+    fn finalize_empty() {
+        let mut buffer = vec![13];
+
+        Prefixer::new(&mut buffer).finalize(&mut buffer).unwrap();
+
+        assert_eq!(buffer.len(), U32_BYTE_LEN + 1);
+        let array: [u8; U32_BYTE_LEN] = buffer[1..].try_into().unwrap();
+        assert_eq!(decode_u32(array), 0);
+    }
+
+    #[test]
+    fn finalize_ok() {
+        let mut buffer = vec![13];
+
+        let prefixer = Prefixer::new(&mut buffer);
+
+        buffer.extend_from_slice(&[0; 42]);
+
+        prefixer.finalize(&mut buffer).unwrap();
+
+        // 1 junk prefix byte, length prefix, 42 bytes of value.
+        assert_eq!(buffer.len(), U32_BYTE_LEN + 43);
+        let prefix = &buffer[1..U32_BYTE_LEN + 1];
+        let array: [u8; U32_BYTE_LEN] = prefix.try_into().unwrap();
+        assert_eq!(decode_u32(array), 42);
+    }
+
+    #[test]
+    #[should_panic]
+    fn finalize_truncated() {
+        let mut buffer = vec![13];
+
+        let prefixer = Prefixer::new(&mut buffer);
+
+        buffer = vec![];
+
+        // Explodes.
+        let _ = prefixer.finalize(&mut buffer);
+    }
+}

src/proto/u32.rs

@@ -0,0 +1,17 @@
+//! This module provides base primitives for encoding and decoding u32 values.
+//!
+//! It mostly centralizes the knowledge that the protocol uses little-endian
+//! representation for u32 values.
+
+/// Length of an encoded 32-bit integer in bytes.
+pub const U32_BYTE_LEN: usize = 4;
+
+/// Returns the byte representatio of the given integer value.
+pub fn encode_u32(value: u32) -> [u8; U32_BYTE_LEN] {
+    value.to_le_bytes()
+}
+
+/// Returns the integer value corresponding to the given bytes.
+pub fn decode_u32(bytes: [u8; U32_BYTE_LEN]) -> u32 {
+    u32::from_le_bytes(bytes)
+}

src/proto/value_codec.rs

@@ -16,17 +16,13 @@
 use std::io;
 use std::net;
 
+use crate::proto::prefix::Prefixer;
+use crate::proto::u32::{decode_u32, encode_u32, U32_BYTE_LEN};
 use encoding::all::WINDOWS_1252;
 use encoding::{DecoderTrap, EncoderTrap, Encoding};
 use std::convert::{TryFrom, TryInto};
 use thiserror::Error;
 
-// Constants
-// ---------
-
-/// Length of an encoded 32-bit integer in bytes.
-pub const U32_BYTE_LEN: usize = 4;
-
 pub trait Decode<T> {
     /// Attempts to decode an instance of `T` from `self`.
     fn decode(&mut self) -> io::Result<T>;
@@ -189,7 +185,7 @@ impl<'a> ValueDecoder<'a> {
         // The conversion from slice to fixed-size array cannot fail, because
         // consume() guarantees that its return value is of size n.
         let array: [u8; U32_BYTE_LEN] = bytes.try_into().unwrap();
-        Ok(u32::from_le_bytes(array))
+        Ok(decode_u32(array))
     }
 
     fn decode_u16(&mut self) -> Result<u16, ValueDecodeError> {
@@ -341,7 +337,7 @@ impl<'a> ValueEncoder<'a> {
 
     /// Encodes the given u32 value into the underlying buffer.
     pub fn encode_u32(&mut self, val: u32) -> Result<(), ValueEncodeError> {
-        self.buffer.extend_from_slice(&val.to_le_bytes());
+        self.buffer.extend_from_slice(&encode_u32(val));
         Ok(())
     }
 
@@ -358,12 +354,8 @@ impl<'a> ValueEncoder<'a> {
 
     /// Encodes the given string into the underlying buffer.
     pub fn encode_string(&mut self, val: &str) -> Result<(), ValueEncodeError> {
-        // Record where we were when we started. This is where we will write
-        // the length prefix once we are done encoding the string. Until then
-        // we do not know how many bytes are needed to encode the string.
-        let prefix_position = self.buffer.len();
-        self.buffer.extend_from_slice(&[0; U32_BYTE_LEN]);
-        let string_position = prefix_position + U32_BYTE_LEN;
+        // Reserve space for the length prefix.
+        let prefixer = Prefixer::new(self.buffer);
 
         // Encode the string. We are quite certain this cannot fail because we
         // use EncoderTrap::Replace to replace any unencodable characters with
@@ -372,22 +364,13 @@ impl<'a> ValueEncoder<'a> {
             .encode_to(val, EncoderTrap::Replace, self.buffer)
             .unwrap();
 
-        // Calculate the length of the string we just encoded.
-        let length = self.buffer.len() - string_position;
-        let length_u32 = match u32::try_from(length) {
-            Ok(value) => value,
-            Err(_) => {
-                return Err(ValueEncodeError::StringTooLong {
-                    string: val.to_string(),
-                    length: length,
-                })
-            }
-        };
-
         // Write the length prefix in the space we initially reserved for it.
-        self.buffer[prefix_position..string_position].copy_from_slice(&length_u32.to_le_bytes());
-
-        Ok(())
+        prefixer
+            .finalize(self.buffer)
+            .map_err(|length| ValueEncodeError::StringTooLong {
+                string: val.to_string(),
+                length: length,
+            })
     }
 
     /// Encodes the given value into the underlying buffer.