diff --git a/src/proto/h2/client.rs b/src/proto/h2/client.rs index eb768baa76..be7592b798 100644 --- a/src/proto/h2/client.rs +++ b/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>, @@ -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, #[pin] send_stream: Option::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, diff --git a/tests/client.rs b/tests/client.rs index 4cb6615657..a5ef009128 100644 --- a/tests/client.rs +++ b/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::::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();