Set content-length for responses to HEAD requests (#734)

* Set `content-length` internally in `Router`

* changelog

* require `B: HttpBody` at the `impl`

* add changelog refs

axum-extra/src/routing/mod.rs

@@ -36,7 +36,7 @@ pub trait RouterExt<B>: sealed::Sealed {
 
 impl<B> RouterExt<B> for Router<B>
 where
-    B: Send + 'static,
+    B: axum::body::HttpBody + Send + 'static,
 {
     fn with<T>(self, routes: T) -> Self
     where

axum-extra/src/routing/resource.rs

@@ -54,7 +54,10 @@ pub struct Resource<B = Body> {
     pub(crate) router: Router<B>,
 }
 
-impl<B: Send + 'static> Resource<B> {
+impl<B> Resource<B>
+where
+    B: axum::body::HttpBody + Send + 'static,
+{
     /// Create a `Resource` with the given name.
     ///
     /// All routes will be nested at `/{resource_name}`.

axum/CHANGELOG.md

@@ -51,13 +51,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
         - `PathRejection`
         - `MatchedPathRejection`
         - `WebSocketUpgradeRejection`
-- **fixed:** Set `Allow` header when responding with `405 Method Not Allowed`
+- **fixed:** Set `Allow` header when responding with `405 Method Not Allowed` ([#733])
+- **fixed:** Correctly set the `Content-Length` header for response to `HEAD`
+  requests ([#734])
 
 [#644]: https://github.com/tokio-rs/axum/pull/644
 [#665]: https://github.com/tokio-rs/axum/pull/665
 [#698]: https://github.com/tokio-rs/axum/pull/698
 [#699]: https://github.com/tokio-rs/axum/pull/699
 [#719]: https://github.com/tokio-rs/axum/pull/719
+[#733]: https://github.com/tokio-rs/axum/pull/733
+[#734]: https://github.com/tokio-rs/axum/pull/734
 
 # 0.4.4 (13. January, 2022)
 

axum/src/routing/future.rs

@@ -1,26 +1,21 @@
 //! Future types.
 
 use crate::response::Response;
-use futures_util::future::Either;
-use std::{convert::Infallible, future::ready};
+use std::convert::Infallible;
 
 pub use super::{into_make_service::IntoMakeServiceFuture, route::RouteFuture};
 
 opaque_future! {
     /// Response future for [`Router`](super::Router).
-    pub type RouterFuture<B> =
-        futures_util::future::Either<
-            RouteFuture<B, Infallible>,
-            std::future::Ready<Result<Response, Infallible>>,
-        >;
+    pub type RouterFuture<B> = RouteFuture<B, Infallible>;
 }
 
 impl<B> RouterFuture<B> {
     pub(super) fn from_future(future: RouteFuture<B, Infallible>) -> Self {
-        Self::new(Either::Left(future))
+        Self::new(future)
     }
 
     pub(super) fn from_response(response: Response) -> Self {
-        Self::new(Either::Right(ready(Ok(response))))
+        Self::new(RouteFuture::from_response(response))
     }
 }

axum/src/routing/method_routing.rs

@@ -960,7 +960,10 @@ where
 
 use crate::routing::future::RouteFuture;
 
-impl<B, E> Service<Request<B>> for MethodRouter<B, E> {
+impl<B, E> Service<Request<B>> for MethodRouter<B, E>
+where
+    B: HttpBody,
+{
     type Response = Response;
     type Error = E;
     type Future = RouteFuture<B, E>;
@@ -980,7 +983,7 @@ impl<B, E> Service<Request<B>> for MethodRouter<B, E> {
             ) => {
                 if $method == Method::$method_variant {
                     if let Some(svc) = $svc {
-                        return RouteFuture::new(svc.0.clone().oneshot($req))
+                        return RouteFuture::from_future(svc.0.clone().oneshot($req))
                             .strip_body($method == Method::HEAD);
                     }
                 }
@@ -1014,13 +1017,14 @@ impl<B, E> Service<Request<B>> for MethodRouter<B, E> {
         call!(req, method, DELETE, delete);
         call!(req, method, TRACE, trace);
 
-        let future =
-            match fallback {
-                Fallback::Default(fallback) => RouteFuture::new(fallback.0.clone().oneshot(req))
-                    .strip_body(method == Method::HEAD),
-                Fallback::Custom(fallback) => RouteFuture::new(fallback.0.clone().oneshot(req))
-                    .strip_body(method == Method::HEAD),
-            };
+        let future = match fallback {
+            Fallback::Default(fallback) => {
+                RouteFuture::from_future(fallback.0.clone().oneshot(req))
+                    .strip_body(method == Method::HEAD)
+            }
+            Fallback::Custom(fallback) => RouteFuture::from_future(fallback.0.clone().oneshot(req))
+                .strip_body(method == Method::HEAD),
+        };
 
         match allow_header {
             AllowHeader::None => future.allow_header(Bytes::new()),

axum/src/routing/mod.rs

@@ -81,7 +81,7 @@ impl<B> Clone for Router<B> {
 
 impl<B> Default for Router<B>
 where
-    B: Send + 'static,
+    B: HttpBody + Send + 'static,
 {
     fn default() -> Self {
         Self::new()
@@ -93,7 +93,7 @@ const NEST_TAIL_PARAM_CAPTURE: &str = "/*axum_nest";
 
 impl<B> Router<B>
 where
-    B: Send + 'static,
+    B: HttpBody + Send + 'static,
 {
     /// Create a new `Router`.
     ///
@@ -461,7 +461,7 @@ where
 
 impl<B> Service<Request<B>> for Router<B>
 where
-    B: Send + 'static,
+    B: HttpBody + Send + 'static,
 {
     type Response = Response;
     type Error = Infallible;

axum/src/routing/route.rs

@@ -1,9 +1,12 @@
 use crate::{
-    body::{boxed, Body, Empty},
+    body::{boxed, Body, Empty, HttpBody},
     response::Response,
 };
 use bytes::Bytes;
-use http::{header, HeaderValue, Request};
+use http::{
+    header::{self, CONTENT_LENGTH},
+    HeaderValue, Request,
+};
 use pin_project_lite::pin_project;
 use std::{
     convert::Infallible,
@@ -46,7 +49,10 @@ impl<ReqBody, E> fmt::Debug for Route<ReqBody, E> {
     }
 }
 
-impl<B, E> Service<Request<B>> for Route<B, E> {
+impl<B, E> Service<Request<B>> for Route<B, E>
+where
+    B: HttpBody,
+{
     type Response = Response;
     type Error = E;
     type Future = RouteFuture<B, E>;
@@ -58,7 +64,7 @@ impl<B, E> Service<Request<B>> for Route<B, E> {
 
     #[inline]
     fn call(&mut self, req: Request<B>) -> Self::Future {
-        RouteFuture::new(self.0.clone().oneshot(req))
+        RouteFuture::from_future(self.0.clone().oneshot(req))
     }
 }
 
@@ -66,21 +72,44 @@ pin_project! {
     /// Response future for [`Route`].
     pub struct RouteFuture<B, E> {
         #[pin]
-        future: Oneshot<
-            BoxCloneService<Request<B>, Response, E>,
-            Request<B>,
-        >,
+        kind: RouteFutureKind<B, E>,
         strip_body: bool,
         allow_header: Option<Bytes>,
     }
 }
 
+pin_project! {
+    #[project = RouteFutureKindProj]
+    enum RouteFutureKind<B, E> {
+        Future {
+            #[pin]
+            future: Oneshot<
+                BoxCloneService<Request<B>, Response, E>,
+                Request<B>,
+            >,
+        },
+        Response {
+            response: Option<Response>,
+        }
+    }
+}
+
 impl<B, E> RouteFuture<B, E> {
-    pub(crate) fn new(
+    pub(crate) fn from_future(
         future: Oneshot<BoxCloneService<Request<B>, Response, E>, Request<B>>,
     ) -> Self {
-        RouteFuture {
-            future,
+        Self {
+            kind: RouteFutureKind::Future { future },
+            strip_body: false,
+            allow_header: None,
+        }
+    }
+
+    pub(super) fn from_response(response: Response) -> Self {
+        Self {
+            kind: RouteFutureKind::Response {
+                response: Some(response),
+            },
             strip_body: false,
             allow_header: None,
         }
@@ -97,38 +126,74 @@ impl<B, E> RouteFuture<B, E> {
     }
 }
 
-impl<B, E> Future for RouteFuture<B, E> {
+impl<B, E> Future for RouteFuture<B, E>
+where
+    B: HttpBody,
+{
     type Output = Result<Response, E>;
 
     #[inline]
-    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let strip_body = self.strip_body;
-        let allow_header = self.allow_header.take();
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.project();
 
-        match self.project().future.poll(cx) {
-            Poll::Ready(Ok(res)) => {
-                let mut res = if strip_body {
-                    res.map(|_| boxed(Empty::new()))
-                } else {
-                    res
-                };
-
-                match allow_header {
-                    Some(allow_header) if !res.headers().contains_key(header::ALLOW) => {
-                        res.headers_mut().insert(
-                            header::ALLOW,
-                            HeaderValue::from_maybe_shared(allow_header)
-                                .expect("invalid `Allow` header"),
-                        );
-                    }
-                    _ => {}
-                }
-
-                Poll::Ready(Ok(res))
+        let mut res = match this.kind.project() {
+            RouteFutureKindProj::Future { future } => match future.poll(cx) {
+                Poll::Ready(Ok(res)) => res,
+                Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
+                Poll::Pending => return Poll::Pending,
+            },
+            RouteFutureKindProj::Response { response } => {
+                response.take().expect("future polled after completion")
             }
-            Poll::Ready(Err(err)) => Poll::Ready(Err(err)),
-            Poll::Pending => Poll::Pending,
+        };
+
+        set_allow_header(&mut res, this.allow_header);
+
+        // make sure to set content-length before removing the body
+        set_content_length(&mut res);
+
+        let res = if *this.strip_body {
+            res.map(|_| boxed(Empty::new()))
+        } else {
+            res
+        };
+
+        Poll::Ready(Ok(res))
+    }
+}
+
+fn set_allow_header<B>(res: &mut Response<B>, allow_header: &mut Option<Bytes>) {
+    match allow_header.take() {
+        Some(allow_header) if !res.headers().contains_key(header::ALLOW) => {
+            res.headers_mut().insert(
+                header::ALLOW,
+                HeaderValue::from_maybe_shared(allow_header).expect("invalid `Allow` header"),
+            );
         }
+        _ => {}
+    }
+}
+
+fn set_content_length<B>(res: &mut Response<B>)
+where
+    B: HttpBody,
+{
+    if res.headers().contains_key(CONTENT_LENGTH) {
+        return;
+    }
+
+    if let Some(size) = res.size_hint().exact() {
+        let header_value = if size == 0 {
+            #[allow(clippy::declare_interior_mutable_const)]
+            const ZERO: HeaderValue = HeaderValue::from_static("0");
+
+            ZERO
+        } else {
+            let mut buffer = itoa::Buffer::new();
+            HeaderValue::from_str(buffer.format(size)).unwrap()
+        };
+
+        res.headers_mut().insert(CONTENT_LENGTH, header_value);
     }
 }
 

axum/src/routing/tests/mod.rs

@@ -618,3 +618,30 @@ async fn merging_routers_with_same_paths_but_different_methods() {
     let body = res.text().await;
     assert_eq!(body, "POST");
 }
+
+#[tokio::test]
+async fn head_content_length_through_hyper_server() {
+    let app = Router::new()
+        .route("/", get(|| async { "foo" }))
+        .route("/json", get(|| async { Json(json!({ "foo": 1 })) }));
+
+    let client = TestClient::new(app);
+
+    let res = client.head("/").send().await;
+    assert_eq!(res.headers()["content-length"], "3");
+    assert!(res.text().await.is_empty());
+
+    let res = client.head("/json").send().await;
+    assert_eq!(res.headers()["content-length"], "9");
+    assert!(res.text().await.is_empty());
+}
+
+#[tokio::test]
+async fn head_content_length_through_hyper_server_that_hits_fallback() {
+    let app = Router::new().fallback((|| async { "foo" }).into_service());
+
+    let client = TestClient::new(app);
+
+    let res = client.head("/").send().await;
+    assert_eq!(res.headers()["content-length"], "3");
+}

axum/src/test_helpers.rs

@@ -59,6 +59,12 @@ impl TestClient {
         }
     }
 
+    pub(crate) fn head(&self, url: &str) -> RequestBuilder {
+        RequestBuilder {
+            builder: self.client.head(format!("http://{}{}", self.addr, url)),
+        }
+    }
+
     pub(crate) fn post(&self, url: &str) -> RequestBuilder {
         RequestBuilder {
             builder: self.client.post(format!("http://{}{}", self.addr, url)),