From c7471d65df73135ee83ca9ca7ef07354d4fa413a Mon Sep 17 00:00:00 2001 From: Titouan Rigoudy Date: Sun, 5 Jul 2020 15:56:18 +0000 Subject: [PATCH] Extract prefix encoding logic into Prefixer. --- src/proto/codec.rs | 3 +- src/proto/mod.rs | 2 + src/proto/prefix.rs | 111 +++++++++++++++++++++++++++++++++++++++ src/proto/u32.rs | 17 ++++++ src/proto/value_codec.rs | 41 +++++---------- 5 files changed, 144 insertions(+), 30 deletions(-) create mode 100644 src/proto/prefix.rs create mode 100644 src/proto/u32.rs diff --git a/src/proto/codec.rs b/src/proto/codec.rs index 825a9ea..d249627 100644 --- a/src/proto/codec.rs +++ b/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 { diff --git a/src/proto/mod.rs b/src/proto/mod.rs index 86295a4..3a90674 100644 --- a/src/proto/mod.rs +++ b/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; diff --git a/src/proto/prefix.rs b/src/proto/prefix.rs new file mode 100644 index 0000000..7113000 --- /dev/null +++ b/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) -> 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) -> 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); + } +} diff --git a/src/proto/u32.rs b/src/proto/u32.rs new file mode 100644 index 0000000..bc6999c --- /dev/null +++ b/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) +} diff --git a/src/proto/value_codec.rs b/src/proto/value_codec.rs index 29cfb8a..dd8b939 100644 --- a/src/proto/value_codec.rs +++ b/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 { /// Attempts to decode an instance of `T` from `self`. fn decode(&mut self) -> io::Result; @@ -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 { @@ -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.