feat(http2): set Content-Length on outgoing messages

hyperium · Jun 10, 2018 · 5b7c7c2 · 5b7c7c2
1 parent f20afba
commit 5b7c7c2
Show file tree

Hide file tree

Showing 5 changed files with 97 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -166,6 +166,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
@@ -353,6 +353,65 @@ mod response_body_lengths {
             expects_con_len: false,
         });
     }
+
+    #[test]
+    fn http2_auto_response_with_known_length() {
+        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.headers().get("content-length").unwrap(), "13");
+                    // TODO: enable this after #1546
+                    let _ = res.body().content_length();
+                    // assert_eq!(res.body().content_length(), Some(13));
+                    Ok(())
+                })
+                .map(|_| ())
+                .map_err(|_e| ())
+        }));
+    }
+
+    #[test]
+    fn http2_auto_response_with_conflicting_lengths() {
+        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| {
+                    assert_eq!(res.headers().get("content-length").unwrap(), "10");
+                    // TODO: enable or remove this after #1546
+                    let _ = res.body().content_length();
+                    // assert_eq!(res.body().content_length(), Some(10));
+                    Ok(())
+                })
+                .map(|_| ())
+                .map_err(|_e| ())
+        }));
+    }
 }
 
 #[test]