From 77ba1de305cb1f3257060219efcee819c1682c6f Mon Sep 17 00:00:00 2001
From: Nikolay Kim <fafhrd91@gmail.com>
Date: Tue, 2 Jan 2018 14:53:51 -0800
Subject: [PATCH] flush encoder

---
 src/encoding.rs | 4 ++++
 src/h1writer.rs | 9 +++------
 src/h2writer.rs | 3 +++
 src/pipeline.rs | 1 +
 4 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/encoding.rs b/src/encoding.rs
index e30ba9f4..aec45929 100644
--- a/src/encoding.rs
+++ b/src/encoding.rs
@@ -569,6 +569,7 @@ impl ContentEncoder {
                 match encoder.finish() {
                     Ok(mut writer) => {
                         writer.encode_eof();
+                        *self = ContentEncoder::Identity(writer);
                         Ok(())
                     },
                     Err(err) => Err(err),
@@ -578,6 +579,7 @@ impl ContentEncoder {
                 match encoder.finish() {
                     Ok(mut writer) => {
                         writer.encode_eof();
+                        *self = ContentEncoder::Identity(writer);
                         Ok(())
                     },
                     Err(err) => Err(err),
@@ -587,6 +589,7 @@ impl ContentEncoder {
                 match encoder.finish() {
                     Ok(mut writer) => {
                         writer.encode_eof();
+                        *self = ContentEncoder::Identity(writer);
                         Ok(())
                     },
                     Err(err) => Err(err),
@@ -594,6 +597,7 @@ impl ContentEncoder {
             },
             ContentEncoder::Identity(mut writer) => {
                 writer.encode_eof();
+                *self = ContentEncoder::Identity(writer);
                 Ok(())
             }
         }
diff --git a/src/h1writer.rs b/src/h1writer.rs
index 011e56ac..5352e743 100644
--- a/src/h1writer.rs
+++ b/src/h1writer.rs
@@ -95,7 +95,6 @@ impl<T: AsyncWrite> H1Writer<T> {
             match self.stream.write(buffer.as_ref()) {
                 Ok(n) => {
                     buffer.split_to(n);
-                    self.written += n as u64;
                 },
                 Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
                     if buffer.len() > MAX_WRITE_BUFFER_SIZE {
@@ -115,11 +114,7 @@ impl<T: AsyncWrite> Writer for H1Writer<T> {
 
     #[cfg_attr(feature = "cargo-clippy", allow(cast_lossless))]
     fn written(&self) -> u64 {
-        if self.written > self.headers_size as u64 {
-            self.written - self.headers_size as u64
-        } else {
-            0
-        }
+        self.written
     }
 
     fn start(&mut self, req: &mut HttpMessage, msg: &mut HttpResponse)
@@ -191,6 +186,7 @@ impl<T: AsyncWrite> Writer for H1Writer<T> {
         }
 
         if let Body::Binary(bytes) = body {
+            self.written = bytes.len() as u64;
             self.encoder.write(bytes.as_ref())?;
         } else {
             msg.replace_body(body);
@@ -199,6 +195,7 @@ impl<T: AsyncWrite> Writer for H1Writer<T> {
     }
 
     fn write(&mut self, payload: &[u8]) -> Result<WriterState, io::Error> {
+        self.written += payload.len() as u64;
         if !self.flags.contains(Flags::DISCONNECTED) {
             if self.flags.contains(Flags::STARTED) {
                 // TODO: add warning, write after EOF
diff --git a/src/h2writer.rs b/src/h2writer.rs
index 51027f7f..0b21421f 100644
--- a/src/h2writer.rs
+++ b/src/h2writer.rs
@@ -167,6 +167,7 @@ impl Writer for H2Writer {
 
         if let Body::Binary(bytes) = body {
             self.flags.insert(Flags::EOF);
+            self.written = bytes.len() as u64;
             self.encoder.write(bytes.as_ref())?;
             if let Some(ref mut stream) = self.stream {
                 stream.reserve_capacity(cmp::min(self.encoder.len(), CHUNK_SIZE));
@@ -179,6 +180,8 @@ impl Writer for H2Writer {
     }
 
     fn write(&mut self, payload: &[u8]) -> Result<WriterState, io::Error> {
+        self.written = payload.len() as u64;
+
         if !self.flags.contains(Flags::DISCONNECTED) {
             if self.flags.contains(Flags::STARTED) {
                 // TODO: add warning, write after EOF
diff --git a/src/pipeline.rs b/src/pipeline.rs
index f5f4799c..e8fbfbe8 100644
--- a/src/pipeline.rs
+++ b/src/pipeline.rs
@@ -677,6 +677,7 @@ impl<S, H> ProcessResponse<S, H> {
         // response is completed
         match self.iostate {
             IOState::Done => {
+                io.write_eof();
                 self.resp.set_response_size(io.written());
                 Ok(FinishingMiddlewares::init(info, self.resp))
             }