From 1879bc49ce3126af4b561bfe535cbec9f7ae5b90 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Thu, 4 Oct 2018 12:46:57 -0700
Subject: [PATCH] codec: Fix panic in `LengthDelimitedCodec::encode` (#682)

Fixes: #681

## Motivation

Currently, a potential panic exists in `LengthDelimitedCodec::encode`.
Writing the length field to the `dst` buffer can exceed the buffer
capacity, as `BufMut::put_uint_{le,be}` doesn't reserve more capacity.

## Solution

This branch adds a call to `dst.reserve` to ensure that there's
sufficient remaining buffer capacity to hold the length field and
the frame, prior to writing the length field. Previously, capacity
was only reserved later in the function, when writing the frame
to the buffer, and we never reserved capacity for the length field.

I've also added a test that reproduces the issue. The test panics on
master, but passes after making this change.

Signed-off-by: Eliza Weisman <eliza@buoyant.io>
---
 src/codec/length_delimited.rs |  4 ++++
 tests/length_delimited.rs     | 18 +++++++++++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

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::<Vec<_>>();
+    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 {