From 673fdb5cb3b16b93e4b5313bdfea952921094e31 Mon Sep 17 00:00:00 2001 From: Eliza Weisman Date: Thu, 30 Aug 2018 14:50:32 -0700 Subject: [PATCH] Refactor `codec::length_delimited` (#575) This patch refactors `length_delimited` to be implemented as a `Codec` and use the default `Framed` wrapper types. The original implementation did not do this in order to support vectored writes in the write half. However, this implementation would be more efficient with small frames anyway. If vectored writes are to be explored in the future, then it should be done holistically. Signed-off-by: Eliza Weisman --- src/length_delimited.rs | 462 +++++-------------------------- src/lib.rs | 25 +- tests/length_delimited.rs | 554 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 640 insertions(+), 401 deletions(-) create mode 100644 tests/length_delimited.rs diff --git a/src/length_delimited.rs b/src/length_delimited.rs index 3f29f8daa..270a3c1f9 100644 --- a/src/length_delimited.rs +++ b/src/length_delimited.rs @@ -1,19 +1,17 @@ -#![allow(deprecated)] +use { + codec::{Decoder, Encoder, FramedRead, FramedWrite, Framed}, + io::{AsyncRead, AsyncWrite}, +}; -use tokio_io::{codec, AsyncRead, AsyncWrite}; - -use bytes::{Buf, BufMut, BytesMut, IntoBuf}; -use bytes::buf::Chain; - -use futures::{Async, AsyncSink, Stream, Sink, StartSend, Poll}; +use bytes::{Buf, BufMut, Bytes, BytesMut, IntoBuf}; use std::{cmp, fmt}; use std::error::Error as StdError; use std::io::{self, Cursor}; -/// Configure length delimited `FramedRead`, `FramedWrite`, and `Framed` values. +/// Configure length delimited `LengthDelimitedCodec`s. /// -/// `Builder` enables constructing configured length delimited framers. Note +/// `Builder` enables constructing configured length delimited codecs. Note /// that not all configuration settings apply to both encoding and decoding. See /// the documentation for specific methods for more detail. #[derive(Debug, Clone, Copy)] @@ -38,33 +36,21 @@ pub struct Builder { length_field_is_big_endian: bool, } -/// Adapts a byte stream into a unified `Stream` and `Sink` that works over -/// entire frame values. -/// -/// See [module level] documentation for more detail. -/// -/// [module level]: index.html -pub struct Framed { - inner: FramedRead>, -} - -/// Adapts a byte stream to a `Stream` yielding entire frame values. -/// -/// See [module level] documentation for more detail. -/// -/// [module level]: index.html -#[derive(Debug)] -pub struct FramedRead { - inner: codec::FramedRead, -} - /// An error when the number of bytes read is more than max frame length. pub struct FrameTooBig { _priv: (), } +/// A codec for frames delimited by a frame head specifying their lengths. +/// +/// This allows the consumer to work with entire frames without having to worry +/// about buffering or other framing logic. +/// +/// See [module level] documentation for more detail. +/// +/// [module level]: index.html #[derive(Debug)] -struct Decoder { +pub struct LengthDelimitedCodec { // Configuration values builder: Builder, @@ -78,114 +64,23 @@ enum DecodeState { Data(usize), } -/// Adapts a byte stream to a `Sink` accepting entire frame values. -/// -/// See [module level] documentation for more detail. -/// -/// [module level]: index.html -pub struct FramedWrite { - // I/O type - inner: T, +// ===== impl LengthDelimitedCodec ====== - // Configuration values - builder: Builder, - - // Current frame being written - frame: Option, B::Buf>>, -} - -// ===== impl Framed ===== - -impl Framed { - /// Creates a new `Framed` with default configuration values. - pub fn new(inner: T) -> Framed { - Builder::new().new_framed(inner) - } -} - -impl Framed { - /// Returns a reference to the underlying I/O stream wrapped by `Framed`. - /// - /// Note that care should be taken to not tamper with the underlying stream - /// of data coming in as it may corrupt the stream of frames otherwise - /// being worked with. - pub fn get_ref(&self) -> &T { - self.inner.get_ref().get_ref() +impl LengthDelimitedCodec { + /// Creates a new `LengthDelimitedCodec` with the default configuration values. + pub fn new() -> Self { + Self { + builder: Builder::new(), + state: DecodeState::Head, + } } - /// Returns a mutable reference to the underlying I/O stream wrapped by - /// `Framed`. - /// - /// Note that care should be taken to not tamper with the underlying stream - /// of data coming in as it may corrupt the stream of frames otherwise being - /// worked with. - pub fn get_mut(&mut self) -> &mut T { - self.inner.get_mut().get_mut() - } - - /// Consumes the `Framed`, returning its underlying I/O stream. - /// - /// Note that care should be taken to not tamper with the underlying stream - /// of data coming in as it may corrupt the stream of frames otherwise being - /// worked with. - pub fn into_inner(self) -> T { - self.inner.into_inner().into_inner() - } -} - -impl Stream for Framed { - type Item = BytesMut; - type Error = io::Error; - - fn poll(&mut self) -> Poll, io::Error> { - self.inner.poll() - } -} - -impl Sink for Framed { - type SinkItem = B; - type SinkError = io::Error; - - fn start_send(&mut self, item: B) -> StartSend { - self.inner.start_send(item) - } - - fn poll_complete(&mut self) -> Poll<(), io::Error> { - self.inner.poll_complete() - } - - fn close(&mut self) -> Poll<(), io::Error> { - self.inner.close() - } -} - -impl fmt::Debug for Framed - where T: fmt::Debug, - B::Buf: fmt::Debug, -{ - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - f.debug_struct("Framed") - .field("inner", &self.inner) - .finish() - } -} - -// ===== impl FramedRead ===== - -impl FramedRead { - /// Creates a new `FramedRead` with default configuration values. - pub fn new(inner: T) -> FramedRead { - Builder::new().new_read(inner) - } -} - -impl FramedRead { /// Returns the current max frame setting /// /// This is the largest size this codec will accept from the wire. Larger /// frames will be rejected. pub fn max_frame_length(&self) -> usize { - self.inner.decoder().builder.max_frame_len + self.builder.max_frame_len } /// Updates the max frame setting. @@ -195,87 +90,9 @@ impl FramedRead { /// size greater than `val` but less than the max frame length in effect /// before calling this function, then the frame will be allowed. pub fn set_max_frame_length(&mut self, val: usize) { - self.inner.decoder_mut().builder.max_frame_length(val); + self.builder.max_frame_length(val); } - /// Returns a reference to the underlying I/O stream wrapped by `FramedRead`. - /// - /// Note that care should be taken to not tamper with the underlying stream - /// of data coming in as it may corrupt the stream of frames otherwise - /// being worked with. - pub fn get_ref(&self) -> &T { - self.inner.get_ref() - } - - /// Returns a mutable reference to the underlying I/O stream wrapped by - /// `FramedRead`. - /// - /// Note that care should be taken to not tamper with the underlying stream - /// of data coming in as it may corrupt the stream of frames otherwise being - /// worked with. - pub fn get_mut(&mut self) -> &mut T { - self.inner.get_mut() - } - - /// Consumes the `FramedRead`, returning its underlying I/O stream. - /// - /// Note that care should be taken to not tamper with the underlying stream - /// of data coming in as it may corrupt the stream of frames otherwise being - /// worked with. - pub fn into_inner(self) -> T { - self.inner.into_inner() - } -} - -impl Stream for FramedRead { - type Item = BytesMut; - type Error = io::Error; - - fn poll(&mut self) -> Poll, io::Error> { - self.inner.poll() - } -} - -impl Sink for FramedRead { - type SinkItem = T::SinkItem; - type SinkError = T::SinkError; - - fn start_send(&mut self, item: T::SinkItem) -> StartSend { - self.inner.start_send(item) - } - - fn poll_complete(&mut self) -> Poll<(), T::SinkError> { - self.inner.poll_complete() - } - - fn close(&mut self) -> Poll<(), T::SinkError> { - self.inner.close() - } -} - -impl io::Write for FramedRead { - fn write(&mut self, src: &[u8]) -> io::Result { - self.inner.get_mut().write(src) - } - - fn flush(&mut self) -> io::Result<()> { - self.inner.get_mut().flush() - } -} - -impl AsyncWrite for FramedRead { - fn shutdown(&mut self) -> Poll<(), io::Error> { - self.inner.get_mut().shutdown() - } - - fn write_buf(&mut self, buf: &mut B) -> Poll { - self.inner.get_mut().write_buf(buf) - } -} - -// ===== impl Decoder ====== - -impl Decoder { fn decode_head(&mut self, src: &mut BytesMut) -> io::Result> { let head_len = self.builder.num_head_bytes(); let field_len = self.builder.length_field_len; @@ -345,7 +162,7 @@ impl Decoder { } } -impl codec::Decoder for Decoder { +impl Decoder for LengthDelimitedCodec { type Item = BytesMut; type Error = io::Error; @@ -378,88 +195,12 @@ impl codec::Decoder for Decoder { } } -// ===== impl FramedWrite ===== +impl Encoder for LengthDelimitedCodec { + type Item = Bytes; + type Error = io::Error; -impl FramedWrite { - /// Creates a new `FramedWrite` with default configuration values. - pub fn new(inner: T) -> FramedWrite { - Builder::new().new_write(inner) - } -} - -impl FramedWrite { - /// Returns the current max frame setting - /// - /// This is the largest size this codec will write to the wire. Larger - /// frames will be rejected. - pub fn max_frame_length(&self) -> usize { - self.builder.max_frame_len - } - - /// Updates the max frame setting. - /// - /// The change takes effect the next time a frame is encoded. In other - /// words, if a frame is currently in process of being encoded with a frame - /// size greater than `val` but less than the max frame length in effect - /// before calling this function, then the frame will be allowed. - pub fn set_max_frame_length(&mut self, val: usize) { - self.builder.max_frame_length(val); - } - - /// Returns a reference to the underlying I/O stream wrapped by - /// `FramedWrite`. - /// - /// Note that care should be taken to not tamper with the underlying stream - /// of data coming in as it may corrupt the stream of frames otherwise - /// being worked with. - pub fn get_ref(&self) -> &T { - &self.inner - } - - /// Returns a mutable reference to the underlying I/O stream wrapped by - /// `FramedWrite`. - /// - /// Note that care should be taken to not tamper with the underlying stream - /// of data coming in as it may corrupt the stream of frames otherwise being - /// worked with. - pub fn get_mut(&mut self) -> &mut T { - &mut self.inner - } - - /// Consumes the `FramedWrite`, returning its underlying I/O stream. - /// - /// Note that care should be taken to not tamper with the underlying stream - /// of data coming in as it may corrupt the stream of frames otherwise being - /// worked with. - pub fn into_inner(self) -> T { - self.inner - } -} - -impl FramedWrite { - // If there is a buffered frame, try to write it to `T` - fn do_write(&mut self) -> Poll<(), io::Error> { - if self.frame.is_none() { - return Ok(Async::Ready(())); - } - - loop { - let frame = self.frame.as_mut().unwrap(); - try_ready!(self.inner.write_buf(frame)); - - if !frame.has_remaining() { - break; - } - } - - self.frame = None; - - Ok(Async::Ready(())) - } - - fn set_frame(&mut self, buf: B::Buf) -> io::Result<()> { - let mut head = BytesMut::with_capacity(8); - let n = buf.remaining(); + fn encode(&mut self, data: Bytes, dst: &mut BytesMut) -> Result<(), io::Error> { + let n = (&data).into_buf().remaining(); if n > self.builder.max_frame_len { return Err(io::Error::new(io::ErrorKind::InvalidInput, FrameTooBig { @@ -474,98 +215,28 @@ impl FramedWrite { n.checked_sub(self.builder.length_adjustment as usize) }; - // Error handling - let n = match n { - Some(n) => n, - None => return Err(io::Error::new(io::ErrorKind::InvalidInput, "provided length would overflow after adjustment")), - }; + let n = n.ok_or_else(|| io::Error::new( + io::ErrorKind::InvalidInput, + "provided length would overflow after adjustment", + ))?; if self.builder.length_field_is_big_endian { - head.put_uint_be(n as u64, self.builder.length_field_len); + dst.put_uint_be(n as u64, self.builder.length_field_len); } else { - head.put_uint_le(n as u64, self.builder.length_field_len); + dst.put_uint_le(n as u64, self.builder.length_field_len); } - debug_assert!(self.frame.is_none()); - - self.frame = Some(head.into_buf().chain(buf)); + // Write the frame to the buffer + dst.extend_from_slice(&data[..]); Ok(()) } } -impl Sink for FramedWrite { - type SinkItem = B; - type SinkError = io::Error; - - fn start_send(&mut self, item: B) -> StartSend { - if !try!(self.do_write()).is_ready() { - return Ok(AsyncSink::NotReady(item)); - } - - try!(self.set_frame(item.into_buf())); - - Ok(AsyncSink::Ready) - } - - fn poll_complete(&mut self) -> Poll<(), io::Error> { - // Write any buffered frame to T - try_ready!(self.do_write()); - - // Try flushing the underlying IO - try_ready!(self.inner.poll_flush()); - - return Ok(Async::Ready(())); - } - - fn close(&mut self) -> Poll<(), io::Error> { - try_ready!(self.poll_complete()); - self.inner.shutdown() - } -} - -impl Stream for FramedWrite { - type Item = T::Item; - type Error = T::Error; - - fn poll(&mut self) -> Poll, T::Error> { - self.inner.poll() - } -} - -impl io::Read for FramedWrite { - fn read(&mut self, dst: &mut [u8]) -> io::Result { - self.get_mut().read(dst) - } -} - -impl AsyncRead for FramedWrite { - fn read_buf(&mut self, buf: &mut B) -> Poll { - self.get_mut().read_buf(buf) - } - - unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool { - self.get_ref().prepare_uninitialized_buffer(buf) - } -} - -impl fmt::Debug for FramedWrite - where T: fmt::Debug, - B::Buf: fmt::Debug, -{ - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - f.debug_struct("FramedWrite") - .field("inner", &self.inner) - .field("builder", &self.builder) - .field("frame", &self.frame) - .finish() - } -} - // ===== impl Builder ===== impl Builder { - /// Creates a new length delimited framer builder with default configuration + /// Creates a new length delimited codec builder with default configuration /// values. /// /// # Examples @@ -813,6 +484,30 @@ impl Builder { self } + /// Create a configured length delimited `LengthDelimitedCodec` + /// + /// # Examples + /// + /// ``` + /// # extern crate tokio; + /// # use tokio::io::AsyncRead; + /// use tokio::codec::length_delimited::Builder; + /// # pub fn main() { + /// Builder::new() + /// .length_field_offset(0) + /// .length_field_length(2) + /// .length_adjustment(0) + /// .num_skip(0) + /// .new_codec(); + /// # } + /// ``` + pub fn new_codec(&self) -> LengthDelimitedCodec { + LengthDelimitedCodec { + builder: *self, + state: DecodeState::Head, + } + } + /// Create a configured length delimited `FramedRead` /// /// # Examples @@ -832,15 +527,10 @@ impl Builder { /// # } /// # pub fn main() {} /// ``` - pub fn new_read(&self, upstream: T) -> FramedRead + pub fn new_read(&self, upstream: T) -> FramedRead where T: AsyncRead, { - FramedRead { - inner: codec::FramedRead::new(upstream, Decoder { - builder: *self, - state: DecodeState::Head, - }), - } + FramedRead::new(upstream, self.new_codec()) } /// Create a configured length delimited `FramedWrite` @@ -854,22 +544,16 @@ impl Builder { /// # use tokio::codec::length_delimited; /// # use bytes::BytesMut; /// # fn write_frame(io: T) { - /// # let _: length_delimited::FramedWrite = /// length_delimited::Builder::new() /// .length_field_length(2) /// .new_write(io); /// # } /// # pub fn main() {} /// ``` - pub fn new_write(&self, inner: T) -> FramedWrite + pub fn new_write(&self, inner: T) -> FramedWrite where T: AsyncWrite, - B: IntoBuf, { - FramedWrite { - inner: inner, - builder: *self, - frame: None, - } + FramedWrite::new(inner, self.new_codec()) } /// Create a configured length delimited `Framed` @@ -883,19 +567,17 @@ impl Builder { /// # use tokio::codec::length_delimited; /// # use bytes::BytesMut; /// # fn write_frame(io: T) { - /// # let _: length_delimited::Framed = + /// # let _ = /// length_delimited::Builder::new() /// .length_field_length(2) /// .new_framed(io); /// # } /// # pub fn main() {} /// ``` - pub fn new_framed(&self, inner: T) -> Framed + pub fn new_framed(&self, inner: T) -> Framed where T: AsyncRead + AsyncWrite, - B: IntoBuf { - let inner = self.new_read(self.new_write(inner)); - Framed { inner: inner } + Framed::new(inner, self.new_codec()) } fn num_head_bytes(&self) -> usize { diff --git a/src/lib.rs b/src/lib.rs index 368fca88b..7bb343062 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -135,19 +135,20 @@ pub mod codec { //! # Getting started //! //! If implementing a protocol from scratch, using length delimited framing - //! is an easy way to get started. [`Framed::new()`] will adapt a - //! full-duplex byte stream with a length delimited framer using default - //! configuration values. + //! is an easy way to get started. [`Codec::new()`] will return a length + //! delimited codec using default configuration values. This can then be + //! used to construct a framer to adapt a full-duplex byte stream into a + //! stream of frames. //! //! ``` //! # extern crate tokio; //! use tokio::io::{AsyncRead, AsyncWrite}; - //! use tokio::codec::length_delimited; + //! use tokio::codec::*; //! //! fn bind_transport(io: T) - //! -> length_delimited::Framed + //! -> Framed //! { - //! length_delimited::Framed::new(io) + //! Framed::new(io, LengthDelimitedCodec::new()) //! } //! # pub fn main() {} //! ``` @@ -170,13 +171,13 @@ pub mod codec { //! # extern crate futures; //! # //! use tokio::io::{AsyncRead, AsyncWrite}; - //! use tokio::codec::length_delimited; - //! use bytes::BytesMut; + //! use tokio::codec::*; + //! use bytes::Bytes; //! use futures::{Sink, Future}; //! //! fn write_frame(io: T) { - //! let mut transport = length_delimited::Framed::new(io); - //! let frame = BytesMut::from("hello world"); + //! let mut transport = Framed::new(io, LengthDelimitedCodec::new()); + //! let frame = Bytes::from("hello world"); //! //! transport.send(frame).wait().unwrap(); //! } @@ -454,7 +455,7 @@ pub mod codec { //! # use tokio::codec::length_delimited; //! # use bytes::BytesMut; //! # fn write_frame(io: T) { - //! # let _: length_delimited::FramedWrite = + //! # let _ = //! length_delimited::Builder::new() //! .length_field_length(2) //! .new_write(io); @@ -478,6 +479,8 @@ pub mod codec { //! [`BytesMut`]: https://docs.rs/bytes/0.4/bytes/struct.BytesMut.html pub use ::length_delimited::*; } + + pub use self::length_delimited::LengthDelimitedCodec; } pub mod io { diff --git a/tests/length_delimited.rs b/tests/length_delimited.rs new file mode 100644 index 000000000..077aaa721 --- /dev/null +++ b/tests/length_delimited.rs @@ -0,0 +1,554 @@ +extern crate tokio; +extern crate futures; +extern crate bytes; + +use tokio::io::{AsyncRead, AsyncWrite}; +use tokio::codec::*; + +use bytes::Bytes; +use futures::{Stream, Sink, Poll}; +use futures::Async::*; + +use std::io; +use std::collections::VecDeque; + +macro_rules! mock { + ($($x:expr,)*) => {{ + let mut v = VecDeque::new(); + v.extend(vec![$($x),*]); + Mock { calls: v } + }}; +} + + +#[test] +fn read_empty_io_yields_nothing() { + let mut io = FramedRead::new(mock!(), LengthDelimitedCodec::new()); + + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_single_frame_one_packet() { + let mut io = FramedRead::new(mock! { + Ok(b"\x00\x00\x00\x09abcdefghi"[..].into()), + }, LengthDelimitedCodec::new()); + + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_single_frame_one_packet_little_endian() { + let mut io = length_delimited::Builder::new() + .little_endian() + .new_read(mock! { + Ok(b"\x09\x00\x00\x00abcdefghi"[..].into()), + }); + + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_single_frame_one_packet_native_endian() { + let data = if cfg!(target_endian = "big") { + b"\x00\x00\x00\x09abcdefghi" + } else { + b"\x09\x00\x00\x00abcdefghi" + }; + let mut io = length_delimited::Builder::new() + .native_endian() + .new_read(mock! { + Ok(data[..].into()), + }); + + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_single_multi_frame_one_packet() { + let mut data: Vec = vec![]; + data.extend_from_slice(b"\x00\x00\x00\x09abcdefghi"); + data.extend_from_slice(b"\x00\x00\x00\x03123"); + data.extend_from_slice(b"\x00\x00\x00\x0bhello world"); + + let mut io = FramedRead::new(mock! { + Ok(data.into()), + }, LengthDelimitedCodec::new()); + + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(Some(b"123"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(Some(b"hello world"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_single_frame_multi_packet() { + let mut io = FramedRead::new(mock! { + Ok(b"\x00\x00"[..].into()), + Ok(b"\x00\x09abc"[..].into()), + Ok(b"defghi"[..].into()), + }, LengthDelimitedCodec::new()); + + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_multi_frame_multi_packet() { + let mut io = FramedRead::new(mock! { + Ok(b"\x00\x00"[..].into()), + Ok(b"\x00\x09abc"[..].into()), + Ok(b"defghi"[..].into()), + Ok(b"\x00\x00\x00\x0312"[..].into()), + Ok(b"3\x00\x00\x00\x0bhello world"[..].into()), + }, LengthDelimitedCodec::new()); + + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(Some(b"123"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(Some(b"hello world"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_single_frame_multi_packet_wait() { + let mut io = FramedRead::new(mock! { + Ok(b"\x00\x00"[..].into()), + Err(would_block()), + Ok(b"\x00\x09abc"[..].into()), + Err(would_block()), + Ok(b"defghi"[..].into()), + Err(would_block()), + }, LengthDelimitedCodec::new()); + + assert_eq!(io.poll().unwrap(), NotReady); + assert_eq!(io.poll().unwrap(), NotReady); + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), NotReady); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_multi_frame_multi_packet_wait() { + let mut io = FramedRead::new(mock! { + Ok(b"\x00\x00"[..].into()), + Err(would_block()), + Ok(b"\x00\x09abc"[..].into()), + Err(would_block()), + Ok(b"defghi"[..].into()), + Err(would_block()), + Ok(b"\x00\x00\x00\x0312"[..].into()), + Err(would_block()), + Ok(b"3\x00\x00\x00\x0bhello world"[..].into()), + Err(would_block()), + }, LengthDelimitedCodec::new()); + + + assert_eq!(io.poll().unwrap(), NotReady); + assert_eq!(io.poll().unwrap(), NotReady); + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), NotReady); + assert_eq!(io.poll().unwrap(), NotReady); + assert_eq!(io.poll().unwrap(), Ready(Some(b"123"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(Some(b"hello world"[..].into()))); + assert_eq!(io.poll().unwrap(), NotReady); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_incomplete_head() { + let mut io = FramedRead::new(mock! { + Ok(b"\x00\x00"[..].into()), + }, LengthDelimitedCodec::new()); + + assert!(io.poll().is_err()); +} + +#[test] +fn read_incomplete_head_multi() { + let mut io = FramedRead::new(mock! { + Err(would_block()), + Ok(b"\x00"[..].into()), + Err(would_block()), + }, LengthDelimitedCodec::new()); + + assert_eq!(io.poll().unwrap(), NotReady); + assert_eq!(io.poll().unwrap(), NotReady); + assert!(io.poll().is_err()); +} + +#[test] +fn read_incomplete_payload() { + let mut io = FramedRead::new(mock! { + Ok(b"\x00\x00\x00\x09ab"[..].into()), + Err(would_block()), + Ok(b"cd"[..].into()), + Err(would_block()), + }, LengthDelimitedCodec::new()); + + assert_eq!(io.poll().unwrap(), NotReady); + assert_eq!(io.poll().unwrap(), NotReady); + assert!(io.poll().is_err()); +} + +#[test] +fn read_max_frame_len() { + let mut io = length_delimited::Builder::new() + .max_frame_length(5) + .new_read(mock! { + Ok(b"\x00\x00\x00\x09abcdefghi"[..].into()), + }); + + assert_eq!(io.poll().unwrap_err().kind(), io::ErrorKind::InvalidData); +} + +#[test] +fn read_update_max_frame_len_at_rest() { + let mut io = length_delimited::Builder::new() + .new_read(mock! { + Ok(b"\x00\x00\x00\x09abcdefghi"[..].into()), + Ok(b"\x00\x00\x00\x09abcdefghi"[..].into()), + }); + + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + io.decoder_mut().set_max_frame_length(5); + assert_eq!(io.poll().unwrap_err().kind(), io::ErrorKind::InvalidData); +} + +#[test] +fn read_update_max_frame_len_in_flight() { + let mut io = length_delimited::Builder::new() + .new_read(mock! { + Ok(b"\x00\x00\x00\x09abcd"[..].into()), + Err(would_block()), + Ok(b"efghi"[..].into()), + Ok(b"\x00\x00\x00\x09abcdefghi"[..].into()), + }); + + assert_eq!(io.poll().unwrap(), NotReady); + io.decoder_mut().set_max_frame_length(5); + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap_err().kind(), io::ErrorKind::InvalidData); +} + +#[test] +fn read_one_byte_length_field() { + let mut io = length_delimited::Builder::new() + .length_field_length(1) + .new_read(mock! { + Ok(b"\x09abcdefghi"[..].into()), + }); + + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_header_offset() { + let mut io = length_delimited::Builder::new() + .length_field_length(2) + .length_field_offset(4) + .new_read(mock! { + Ok(b"zzzz\x00\x09abcdefghi"[..].into()), + }); + + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_single_multi_frame_one_packet_skip_none_adjusted() { + let mut data: Vec = vec![]; + data.extend_from_slice(b"xx\x00\x09abcdefghi"); + data.extend_from_slice(b"yy\x00\x03123"); + data.extend_from_slice(b"zz\x00\x0bhello world"); + + let mut io = length_delimited::Builder::new() + .length_field_length(2) + .length_field_offset(2) + .num_skip(0) + .length_adjustment(4) + .new_read(mock! { + Ok(data.into()), + }); + + assert_eq!(io.poll().unwrap(), Ready(Some(b"xx\x00\x09abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(Some(b"yy\x00\x03123"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(Some(b"zz\x00\x0bhello world"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn read_single_multi_frame_one_packet_length_includes_head() { + let mut data: Vec = vec![]; + data.extend_from_slice(b"\x00\x0babcdefghi"); + data.extend_from_slice(b"\x00\x05123"); + data.extend_from_slice(b"\x00\x0dhello world"); + + let mut io = length_delimited::Builder::new() + .length_field_length(2) + .length_adjustment(-2) + .new_read(mock! { + Ok(data.into()), + }); + + assert_eq!(io.poll().unwrap(), Ready(Some(b"abcdefghi"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(Some(b"123"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(Some(b"hello world"[..].into()))); + assert_eq!(io.poll().unwrap(), Ready(None)); +} + +#[test] +fn write_single_frame_length_adjusted() { + let mut io = length_delimited::Builder::new() + .length_adjustment(-2) + .new_write(mock! { + Ok(b"\x00\x00\x00\x0b"[..].into()), + Ok(b"abcdefghi"[..].into()), + Ok(Flush), + }); + assert!(io.start_send(Bytes::from("abcdefghi")).unwrap().is_ready()); + assert!(io.poll_complete().unwrap().is_ready()); + assert!(io.get_ref().calls.is_empty()); +} + +#[test] +fn write_nothing_yields_nothing() { + let mut io = FramedWrite::new( + mock!(), + LengthDelimitedCodec::new() + ); + assert!(io.poll_complete().unwrap().is_ready()); +} + +#[test] +fn write_single_frame_one_packet() { + let mut io = FramedWrite::new(mock! { + Ok(b"\x00\x00\x00\x09"[..].into()), + Ok(b"abcdefghi"[..].into()), + Ok(Flush), + }, LengthDelimitedCodec::new()); + + assert!(io.start_send(Bytes::from("abcdefghi")).unwrap().is_ready()); + assert!(io.poll_complete().unwrap().is_ready()); + assert!(io.get_ref().calls.is_empty()); +} + +#[test] +fn write_single_multi_frame_one_packet() { + let mut io = FramedWrite::new(mock! { + Ok(b"\x00\x00\x00\x09"[..].into()), + Ok(b"abcdefghi"[..].into()), + Ok(b"\x00\x00\x00\x03"[..].into()), + Ok(b"123"[..].into()), + Ok(b"\x00\x00\x00\x0b"[..].into()), + Ok(b"hello world"[..].into()), + Ok(Flush), + }, LengthDelimitedCodec::new()); + + assert!(io.start_send(Bytes::from("abcdefghi")).unwrap().is_ready()); + assert!(io.start_send(Bytes::from("123")).unwrap().is_ready()); + assert!(io.start_send(Bytes::from("hello world")).unwrap().is_ready()); + assert!(io.poll_complete().unwrap().is_ready()); + assert!(io.get_ref().calls.is_empty()); +} + +#[test] +fn write_single_multi_frame_multi_packet() { + let mut io = FramedWrite::new(mock! { + Ok(b"\x00\x00\x00\x09"[..].into()), + Ok(b"abcdefghi"[..].into()), + Ok(Flush), + Ok(b"\x00\x00\x00\x03"[..].into()), + Ok(b"123"[..].into()), + Ok(Flush), + Ok(b"\x00\x00\x00\x0b"[..].into()), + Ok(b"hello world"[..].into()), + Ok(Flush), + }, LengthDelimitedCodec::new()); + + assert!(io.start_send(Bytes::from("abcdefghi")).unwrap().is_ready()); + assert!(io.poll_complete().unwrap().is_ready()); + assert!(io.start_send(Bytes::from("123")).unwrap().is_ready()); + assert!(io.poll_complete().unwrap().is_ready()); + assert!(io.start_send(Bytes::from("hello world")).unwrap().is_ready()); + assert!(io.poll_complete().unwrap().is_ready()); + assert!(io.get_ref().calls.is_empty()); +} + +#[test] +fn write_single_frame_would_block() { + let mut io = FramedWrite::new(mock! { + Err(would_block()), + Ok(b"\x00\x00"[..].into()), + Err(would_block()), + Ok(b"\x00\x09"[..].into()), + Ok(b"abcdefghi"[..].into()), + Ok(Flush), + }, LengthDelimitedCodec::new()); + + assert!(io.start_send(Bytes::from("abcdefghi")).unwrap().is_ready()); + assert!(!io.poll_complete().unwrap().is_ready()); + assert!(!io.poll_complete().unwrap().is_ready()); + assert!(io.poll_complete().unwrap().is_ready()); + + assert!(io.get_ref().calls.is_empty()); +} + +#[test] +fn write_single_frame_little_endian() { + let mut io = length_delimited::Builder::new() + .little_endian() + .new_write(mock! { + Ok(b"\x09\x00\x00\x00"[..].into()), + Ok(b"abcdefghi"[..].into()), + Ok(Flush), + }); + + assert!(io.start_send(Bytes::from("abcdefghi")).unwrap().is_ready()); + assert!(io.poll_complete().unwrap().is_ready()); + assert!(io.get_ref().calls.is_empty()); +} + + +#[test] +fn write_single_frame_with_short_length_field() { + let mut io = length_delimited::Builder::new() + .length_field_length(1) + .new_write(mock! { + Ok(b"\x09"[..].into()), + Ok(b"abcdefghi"[..].into()), + Ok(Flush), + }); + + assert!(io.start_send(Bytes::from("abcdefghi")).unwrap().is_ready()); + assert!(io.poll_complete().unwrap().is_ready()); + assert!(io.get_ref().calls.is_empty()); +} + +#[test] +fn write_max_frame_len() { + let mut io = length_delimited::Builder::new() + .max_frame_length(5) + .new_write(mock! { }); + + assert_eq!(io.start_send(Bytes::from("abcdef")).unwrap_err().kind(), io::ErrorKind::InvalidInput); + assert!(io.get_ref().calls.is_empty()); +} + +#[test] +fn write_update_max_frame_len_at_rest() { + let mut io = length_delimited::Builder::new() + .new_write(mock! { + Ok(b"\x00\x00\x00\x06"[..].into()), + Ok(b"abcdef"[..].into()), + Ok(Flush), + }); + + assert!(io.start_send(Bytes::from("abcdef")).unwrap().is_ready()); + assert!(io.poll_complete().unwrap().is_ready()); + io.encoder_mut().set_max_frame_length(5); + assert_eq!(io.start_send(Bytes::from("abcdef")).unwrap_err().kind(), io::ErrorKind::InvalidInput); + assert!(io.get_ref().calls.is_empty()); +} + +#[test] +fn write_update_max_frame_len_in_flight() { + let mut io = length_delimited::Builder::new() + .new_write(mock! { + Ok(b"\x00\x00\x00\x06"[..].into()), + Ok(b"ab"[..].into()), + Err(would_block()), + Ok(b"cdef"[..].into()), + Ok(Flush), + }); + + assert!(io.start_send(Bytes::from("abcdef")).unwrap().is_ready()); + assert!(!io.poll_complete().unwrap().is_ready()); + io.encoder_mut().set_max_frame_length(5); + assert!(io.poll_complete().unwrap().is_ready()); + assert_eq!(io.start_send(Bytes::from("abcdef")).unwrap_err().kind(), io::ErrorKind::InvalidInput); + assert!(io.get_ref().calls.is_empty()); +} + +// ===== Test utils ===== + +fn would_block() -> io::Error { + io::Error::new(io::ErrorKind::WouldBlock, "would block") +} + +struct Mock { + calls: VecDeque>, +} + +enum Op { + Data(Vec), + Flush, +} + +use self::Op::*; + +impl io::Read for Mock { + fn read(&mut self, dst: &mut [u8]) -> io::Result { + match self.calls.pop_front() { + Some(Ok(Op::Data(data))) => { + debug_assert!(dst.len() >= data.len()); + dst[..data.len()].copy_from_slice(&data[..]); + Ok(data.len()) + } + Some(Ok(_)) => panic!(), + Some(Err(e)) => Err(e), + None => Ok(0), + } + } +} + +impl AsyncRead for Mock { +} + +impl io::Write for Mock { + fn write(&mut self, src: &[u8]) -> io::Result { + match self.calls.pop_front() { + Some(Ok(Op::Data(data))) => { + let len = data.len(); + assert!(src.len() >= len, "expect={:?}; actual={:?}", data, src); + assert_eq!(&data[..], &src[..len]); + Ok(len) + } + Some(Ok(_)) => panic!(), + Some(Err(e)) => Err(e), + None => Ok(0), + } + } + + fn flush(&mut self) -> io::Result<()> { + match self.calls.pop_front() { + Some(Ok(Op::Flush)) => { + Ok(()) + } + Some(Ok(_)) => panic!(), + Some(Err(e)) => Err(e), + None => Ok(()), + } + } +} + +impl AsyncWrite for Mock { + fn shutdown(&mut self) -> Poll<(), io::Error> { + Ok(Ready(())) + } +} + +impl<'a> From<&'a [u8]> for Op { + fn from(src: &'a [u8]) -> Op { + Op::Data(src.into()) + } +} + +impl From> for Op { + fn from(src: Vec) -> Op { + Op::Data(src) + } +}