From 58cc0dfbc5fad1cab85e61724c98dd4655bb43d8 Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Sun, 15 Apr 2018 10:22:09 -0700
Subject: [PATCH] Fix Client Request with custom Body Stream halting on certain
 size requests #176

---
 CHANGES.md             |  2 ++
 src/client/pipeline.rs | 13 ++++++++-----
 src/client/writer.rs   |  4 ++++
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 453b7b2f..1fae933b 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -9,6 +9,8 @@
 
 * Fix StaticFiles does not support percent encoded paths #177
 
+* Fix Client Request with custom Body Stream halting on certain size requests #176
+
 
 ## 0.5.1 (2018-04-12)
 
diff --git a/src/client/pipeline.rs b/src/client/pipeline.rs
index 05fcf812..1db68b43 100644
--- a/src/client/pipeline.rs
+++ b/src/client/pipeline.rs
@@ -184,6 +184,7 @@ impl Future for SendRequest {
                         parser: Some(HttpResponseParser::default()),
                         parser_buf: BytesMut::new(),
                         disconnected: false,
+                        body_completed: false,
                         drain: None,
                         decompress: None,
                         should_decompress: self.req.response_decompress(),
@@ -217,6 +218,7 @@ impl Future for SendRequest {
 
 pub(crate) struct Pipeline {
     body: IoBody,
+    body_completed: bool,
     conn: Option<Connection>,
     writer: HttpClientWriter,
     parser: Option<HttpResponseParser>,
@@ -394,7 +396,7 @@ impl Pipeline {
                     IoBody::Payload(mut body) => match body.poll()? {
                         Async::Ready(None) => {
                             self.writer.write_eof()?;
-                            self.disconnected = true;
+                            self.body_completed = true;
                             break;
                         }
                         Async::Ready(Some(chunk)) => {
@@ -421,8 +423,7 @@ impl Pipeline {
                                 for frame in vec {
                                     match frame {
                                         Frame::Chunk(None) => {
-                                            // info.context = Some(ctx);
-                                            self.disconnected = true;
+                                            self.body_completed = true;
                                             self.writer.write_eof()?;
                                             break 'outter;
                                         }
@@ -451,7 +452,7 @@ impl Pipeline {
                         }
                     }
                     IoBody::Done => {
-                        self.disconnected = true;
+                        self.body_completed = true;
                         done = true;
                         break;
                     }
@@ -472,7 +473,9 @@ impl Pipeline {
             .poll_completed(self.conn.as_mut().unwrap(), false)
         {
             Ok(Async::Ready(_)) => {
-                if self.disconnected {
+                if self.disconnected
+                    || (self.body_completed && self.writer.is_completed())
+                {
                     self.write_state = RunningState::Done;
                 } else {
                     self.write_state.resume();
diff --git a/src/client/writer.rs b/src/client/writer.rs
index 8d554b9b..d0300d59 100644
--- a/src/client/writer.rs
+++ b/src/client/writer.rs
@@ -61,6 +61,10 @@ impl HttpClientWriter {
         self.buffer.take();
     }
 
+    pub fn is_completed(&mut self) -> bool {
+        self.buffer.is_empty()
+    }
+
     // pub fn keepalive(&self) -> bool {
     // self.flags.contains(Flags::KEEPALIVE) &&
     // !self.flags.contains(Flags::UPGRADE) }