fix(http2): report zero body size for HEAD responses

src/proto/h2/client.rs

@@ -413,6 +413,7 @@ where
     B: Body,
 {
     is_connect: bool,
+    is_head: bool,
     eos: bool,
     fut: ResponseFuture,
     body_tx: SendStream<SendBuf<B::Data>>,
@@ -567,6 +568,7 @@ where
             send_when: SendWhen {
                 when: ResponseFutMap {
                     fut: f.fut,
+                    is_head: f.is_head,
                     ping: Some(ping),
                     send_stream: Some(send_stream),
                     exec: self.executor.clone(),
@@ -586,6 +588,7 @@ pin_project! {
     {
         #[pin]
         fut: ResponseFuture,
+        is_head: bool,
         ping: Option<Recorder>,
         #[pin]
         send_stream: Option<Option<SendStream<SendBuf<<B as Body>::Data>>>>,
@@ -623,7 +626,12 @@ where
                 // record that we got the response headers
                 ping.record_non_data();
 
-                let content_length = headers::content_length_parse_all(res.headers());
+                let content_length = if *this.is_head {
+                    // HEAD Content-Length describes the corresponding GET response body.
+                    Some(0)
+                } else {
+                    headers::content_length_parse_all(res.headers())
+                };
                 if let (Some(mut send_stream), StatusCode::OK) = (send_stream, res.status()) {
                     if content_length.map_or(false, |len| len != 0) {
                         warn!("h2 connect response with non-zero body not supported");
@@ -714,6 +722,7 @@ where
                     }
 
                     let is_connect = req.method() == Method::CONNECT;
+                    let is_head = req.method() == Method::HEAD;
                     let eos = body.is_end_stream();
 
                     if is_connect
@@ -746,6 +755,7 @@ where
 
                     let f = FutCtx {
                         is_connect,
+                        is_head,
                         eos,
                         fut,
                         body_tx,

tests/client.rs

@@ -2916,6 +2916,98 @@ mod conn {
         assert_eq!(trailers.get("grpc").unwrap(), "0");
     }
 
+    #[tokio::test]
+    async fn http2_head_response_size_hint() {
+        tokio::time::timeout(Duration::from_secs(5), async {
+            for (method, end_stream, trailers, content_length) in [
+                (Method::HEAD, true, false, Some("5")),
+                (Method::HEAD, false, false, Some("5")),
+                (Method::HEAD, false, true, Some("5")),
+                (Method::HEAD, false, false, None),
+                (Method::GET, false, false, Some("5")),
+            ] {
+                let (client_io, server_io) = tokio::io::duplex(4096);
+                let (finish_tx, finish_rx) = oneshot::channel();
+                let expected_body = if method == Method::HEAD {
+                    Bytes::new()
+                } else {
+                    Bytes::from_static(b"hello")
+                };
+                let response_body = expected_body.clone();
+                let expected_method = method.clone();
+                let server = tokio::spawn(async move {
+                    let mut h2 = h2::server::handshake(server_io).await.unwrap();
+                    let (req, mut respond) = h2.accept().await.unwrap().unwrap();
+                    assert_eq!(req.method(), expected_method);
+                    drop(req);
+
+                    let mut response = Response::builder();
+                    if let Some(value) = content_length {
+                        response = response.header("content-length", value);
+                    }
+                    let mut stream = respond
+                        .send_response(response.body(()).unwrap(), end_stream)
+                        .unwrap();
+                    let finish = async move {
+                        finish_rx.await.unwrap();
+                        if !end_stream {
+                            if trailers {
+                                let mut headers = hyper::HeaderMap::new();
+                                headers.insert("x-finished", "yes".parse().unwrap());
+                                stream.send_trailers(headers).unwrap();
+                            } else {
+                                stream.send_data(response_body, true).unwrap();
+                            }
+                        }
+                    };
+                    let ((), result) = future::join(finish, poll_fn(|cx| h2.poll_closed(cx))).await;
+                    result.unwrap();
+                });
+
+                let (mut client, connection) = conn::http2::Builder::new(TokioExecutor)
+                    .handshake(TokioIo::new(client_io))
+                    .await
+                    .unwrap();
+                let connection = tokio::spawn(connection);
+                let request = Request::builder()
+                    .method(method)
+                    .uri("https://localhost/")
+                    .body(Empty::<Bytes>::new())
+                    .unwrap();
+                let response = client.send_request(request).await.unwrap();
+                assert_eq!(response.status(), StatusCode::OK);
+                assert_eq!(
+                    response
+                        .headers()
+                        .get("content-length")
+                        .map(|v| v.to_str().unwrap()),
+                    content_length,
+                );
+                let mut body = response.into_body();
+                assert_eq!(body.size_hint().exact(), Some(expected_body.len() as u64));
+                assert_eq!(body.is_end_stream(), end_stream);
+
+                finish_tx.send(()).unwrap();
+                let collected = (&mut body).collect().await.unwrap();
+                if trailers {
+                    assert_eq!(collected.trailers().unwrap()["x-finished"], "yes");
+                } else {
+                    assert!(collected.trailers().is_none());
+                }
+                assert_eq!(collected.to_bytes(), expected_body);
+                assert_eq!(body.size_hint().exact(), Some(0));
+                assert!(body.is_end_stream());
+                drop(body);
+
+                drop(client);
+                connection.await.unwrap().unwrap();
+                server.await.unwrap();
+            }
+        })
+        .await
+        .expect("response bodies should complete");
+    }
+
     #[tokio::test]
     async fn h2_connect() {
         let (client_io, server_io, _) = setup_duplex_test_server();