From bc91bc5022ae76b43ece9f5dd5df89f517af56c4 Mon Sep 17 00:00:00 2001
From: Eliza Weisman <eliza@buoyant.io>
Date: Fri, 31 Aug 2018 03:31:43 -0700
Subject: [PATCH] Fix non-terminating loop in
 `tokio_io::length_delimited::FramedWrite` (#576)

* tokio-io: fix non-terminating loop in length_delimited::FramedWrite (#497)
---
 tests/length_delimited.rs          | 10 ++++++++++
 tokio-io/src/length_delimited.rs   |  7 ++++++-
 tokio-io/tests/length_delimited.rs | 10 ++++++++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/tests/length_delimited.rs b/tests/length_delimited.rs
index 077aaa721..318f35ef3 100644
--- a/tests/length_delimited.rs
+++ b/tests/length_delimited.rs
@@ -473,6 +473,16 @@ fn write_update_max_frame_len_in_flight() {
     assert!(io.get_ref().calls.is_empty());
 }
 
+#[test]
+fn write_zero() {
+    let mut io = length_delimited::Builder::new()
+        .new_write(mock! { });
+
+    assert!(io.start_send(Bytes::from("abcdef")).unwrap().is_ready());
+    assert_eq!(io.poll_complete().unwrap_err().kind(), io::ErrorKind::WriteZero);
+    assert!(io.get_ref().calls.is_empty());
+}
+
 // ===== Test utils =====
 
 fn would_block() -> io::Error {
diff --git a/tokio-io/src/length_delimited.rs b/tokio-io/src/length_delimited.rs
index 52ebdfaa1..bf55bc881 100644
--- a/tokio-io/src/length_delimited.rs
+++ b/tokio-io/src/length_delimited.rs
@@ -455,7 +455,12 @@ impl<T: AsyncWrite, B: IntoBuf> FramedWrite<T, B> {
 
         loop {
             let frame = self.frame.as_mut().unwrap();
-            try_ready!(self.inner.write_buf(frame));
+            if try_ready!(self.inner.write_buf(frame)) == 0 {
+                return Err(io::Error::new(
+                    io::ErrorKind::WriteZero,
+                    "failed to write frame to transport",
+                ));
+            }
 
             if !frame.has_remaining() {
                 break;
diff --git a/tokio-io/tests/length_delimited.rs b/tokio-io/tests/length_delimited.rs
index b51cda2c1..985a08361 100644
--- a/tokio-io/tests/length_delimited.rs
+++ b/tokio-io/tests/length_delimited.rs
@@ -433,6 +433,16 @@ fn write_max_frame_len() {
     assert!(io.get_ref().calls.is_empty());
 }
 
+#[test]
+fn write_zero() {
+    let mut io = Builder::new()
+        .new_write(mock! { });
+
+    assert!(io.start_send("abcdef").unwrap().is_ready());
+    assert_eq!(io.poll_complete().unwrap_err().kind(), io::ErrorKind::WriteZero);
+    assert!(io.get_ref().calls.is_empty());
+}
+
 #[test]
 fn write_update_max_frame_len_at_rest() {
     let mut io = Builder::new()