From 81a59fe6a49ac6661b1332f5eac682ad0b0a5c79 Mon Sep 17 00:00:00 2001
From: Laurentiu Nicola <lnicola@dend.ro>
Date: Sat, 9 Jun 2018 00:20:33 +0300
Subject: [PATCH] feat(http2): set Content-Length on outgoing messages

---
 src/headers.rs         |  7 +++++
 src/proto/h2/client.rs |  4 +++
 src/proto/h2/server.rs |  4 +++
 tests/integration.rs   | 23 ++++++++++++++++
 tests/server.rs        | 60 ++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 98 insertions(+)

diff --git a/src/headers.rs b/src/headers.rs
index 70b963f69b..593f980144 100644
--- a/src/headers.rs
+++ b/src/headers.rs
@@ -78,6 +78,13 @@ pub fn content_length_value(len: u64) -> HeaderValue {
     }
 }
 
+pub fn set_content_length_if_missing(headers: &mut HeaderMap, len: u64) {
+    headers
+        .entry(CONTENT_LENGTH)
+        .unwrap()
+        .or_insert(content_length_value(len));
+}
+
 pub fn transfer_encoding_is_chunked(headers: &HeaderMap) -> bool {
     is_chunked(headers.get_all(TRANSFER_ENCODING).into_iter())
 }
diff --git a/src/proto/h2/client.rs b/src/proto/h2/client.rs
index 0915eb37be..6a1b65e84c 100644
--- a/src/proto/h2/client.rs
+++ b/src/proto/h2/client.rs
@@ -7,6 +7,7 @@ use tokio_io::{AsyncRead, AsyncWrite};
 
 use body::Payload;
 use ::common::{Exec, Never};
+use headers;
 use super::{PipeToSendStream, SendBuf};
 use ::{Body, Request, Response};
 
@@ -106,6 +107,9 @@ where
                             let (head, body) = req.into_parts();
                             let mut req = ::http::Request::from_parts(head, ());
                             super::strip_connection_headers(req.headers_mut());
+                            if let Some(len) = body.content_length() {
+                                headers::set_content_length_if_missing(req.headers_mut(), len);
+                            }
                             let eos = body.is_end_stream();
                             let (fut, body_tx) = match tx.send_request(req, eos) {
                                 Ok(ok) => ok,
diff --git a/src/proto/h2/server.rs b/src/proto/h2/server.rs
index fa6e7feb82..ad23c0c6ec 100644
--- a/src/proto/h2/server.rs
+++ b/src/proto/h2/server.rs
@@ -5,6 +5,7 @@ use tokio_io::{AsyncRead, AsyncWrite};
 
 use ::body::Payload;
 use ::common::Exec;
+use ::headers;
 use ::service::Service;
 use super::{PipeToSendStream, SendBuf};
 
@@ -171,6 +172,9 @@ where
                     let (head, body) = res.into_parts();
                     let mut res = ::http::Response::from_parts(head, ());
                     super::strip_connection_headers(res.headers_mut());
+                    if let Some(len) = body.content_length() {
+                        headers::set_content_length_if_missing(res.headers_mut(), len);
+                    }
                     macro_rules! reply {
                         ($eos:expr) => ({
                             match self.reply.send_response(res, $eos) {
diff --git a/tests/integration.rs b/tests/integration.rs
index 16592877d6..43ef036e92 100644
--- a/tests/integration.rs
+++ b/tests/integration.rs
@@ -88,6 +88,29 @@ t! {
             ;
 }
 
+t! {
+    post_outgoing_length,
+    client:
+        request:
+            method: "POST",
+            uri: "/hello",
+            body: "hello, world!",
+            ;
+        response:
+            ;
+    server:
+        request:
+            method: "POST",
+            uri: "/hello",
+            headers: {
+                "content-length" => "13",
+            },
+            body: "hello, world!",
+            ;
+        response:
+            ;
+}
+
 t! {
     post_chunked,
     client:
diff --git a/tests/server.rs b/tests/server.rs
index 000c709841..f430da13e7 100644
--- a/tests/server.rs
+++ b/tests/server.rs
@@ -353,6 +353,66 @@ mod response_body_lengths {
             expects_con_len: false,
         });
     }
+
+    #[test]
+    #[ignore]
+    fn http2_auto_response_with_known_length() {
+        // TODO: check that this really works after #1546
+        use hyper::body::Payload;
+
+        let server = serve();
+        let addr_str = format!("http://{}", server.addr());
+        server.reply().body("Hello, World!");
+
+        hyper::rt::run(hyper::rt::lazy(move || {
+            let client: Client<_, hyper::Body> = Client::builder().http2_only(true).build_http();
+            let uri = addr_str
+                .parse::<hyper::Uri>()
+                .expect("server addr should parse");
+
+            client
+                .get(uri)
+                .and_then(|res| {
+                    assert_eq!(res.body().content_length(), Some(13));
+                    Ok(())
+                })
+                .map(|_| ())
+                .map_err(|_e| ())
+        }));
+    }
+
+    #[test]
+    #[ignore]
+    fn http2_auto_response_with_conflicting_lengths() {
+        // TODO: check that this really works after #1546
+        //       if it doesn't, remove it.
+        use hyper::body::Payload;
+
+        let server = serve();
+        let addr_str = format!("http://{}", server.addr());
+        server
+            .reply()
+            .header("content-length", "10")
+            .body("Hello, World!");
+
+        hyper::rt::run(hyper::rt::lazy(move || {
+            let client: Client<_, hyper::Body> = Client::builder().http2_only(true).build_http();
+            let uri = addr_str
+                .parse::<hyper::Uri>()
+                .expect("server addr should parse");
+
+            client
+                .get(uri)
+                .and_then(|res| {
+                    // TODO: remove eprintln
+                    eprintln!("{:?}", res);
+                    assert_eq!(res.body().content_length(), Some(10));
+                    Ok(())
+                })
+                .map(|_| ())
+                .map_err(|_e| ())
+        }));
+    }
 }
 
 #[test]