diff --git a/src/codec/length_delimited.rs b/src/codec/length_delimited.rs index 54ec202bb..b47886513 100644 --- a/src/codec/length_delimited.rs +++ b/src/codec/length_delimited.rs @@ -578,6 +578,10 @@ impl Encoder for LengthDelimitedCodec { "provided length would overflow after adjustment", ))?; + // Reserve capacity in the destination buffer to fit the frame and + // length field (plus adjustment). + dst.reserve(self.builder.length_field_len + n); + if self.builder.length_field_is_big_endian { dst.put_uint_be(n as u64, self.builder.length_field_len); } else { diff --git a/tests/length_delimited.rs b/tests/length_delimited.rs index 318f35ef3..4e118d379 100644 --- a/tests/length_delimited.rs +++ b/tests/length_delimited.rs @@ -5,7 +5,7 @@ extern crate bytes; use tokio::io::{AsyncRead, AsyncWrite}; use tokio::codec::*; -use bytes::Bytes; +use bytes::{Bytes, BytesMut, BufMut}; use futures::{Stream, Sink, Poll}; use futures::Async::*; @@ -483,6 +483,22 @@ fn write_zero() { assert!(io.get_ref().calls.is_empty()); } +#[test] +fn encode_overflow() { + // Test reproducing tokio-rs/tokio#681. + let mut codec = length_delimited::Builder::new().new_codec(); + let mut buf = BytesMut::with_capacity(1024); + + // Put some data into the buffer without resizing it to hold more. + let some_as = std::iter::repeat(b'a') + .take(1024) + .collect::>(); + buf.put_slice(&some_as[..]); + + // Trying to encode the length header should resize the buffer if it won't fit. + codec.encode(Bytes::from("hello"), &mut buf).unwrap(); +} + // ===== Test utils ===== fn would_block() -> io::Error {