Correctly handle HEAD requests (#129)

src/handler/future.rs

@@ -1,7 +1,8 @@
 //! Handler future types.
 
-use crate::body::BoxBody;
-use http::{Request, Response};
+use crate::body::{box_body, BoxBody};
+use http::{Method, Request, Response};
+use http_body::Empty;
 use pin_project_lite::pin_project;
 use std::{
     convert::Infallible,
@@ -27,6 +28,7 @@ pin_project! {
     {
         #[pin]
         pub(super) inner: crate::routing::future::RouteFuture<S, F, B>,
+        pub(super) req_method: Method,
     }
 }
 
@@ -38,6 +40,13 @@ where
     type Output = Result<Response<BoxBody>, S::Error>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        self.project().inner.poll(cx)
+        let this = self.project();
+        let response = futures_util::ready!(this.inner.poll(cx))?;
+        if this.req_method == &Method::HEAD {
+            let response = response.map(|_| box_body(Empty::new()));
+            Poll::Ready(Ok(response))
+        } else {
+            Poll::Ready(Ok(response))
+        }
     }
 }

src/handler/mod.rs

@@ -79,11 +79,15 @@ where
 /// # axum::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap();
 /// # };
 /// ```
+///
+/// Note that `get` routes will also be called for `HEAD` requests but will have
+/// the response body removed. Make sure to add explicit `HEAD` routes
+/// afterwards.
 pub fn get<H, B, T>(handler: H) -> OnMethod<IntoService<H, B, T>, EmptyRouter>
 where
     H: Handler<B, T>,
 {
-    on(MethodFilter::GET, handler)
+    on(MethodFilter::GET | MethodFilter::HEAD, handler)
 }
 
 /// Route `HEAD` requests to the given handler.
@@ -508,11 +512,15 @@ impl<S, F> OnMethod<S, F> {
     /// # axum::Server::bind(&"".parse().unwrap()).serve(app.into_make_service()).await.unwrap();
     /// # };
     /// ```
+    ///
+    /// Note that `get` routes will also be called for `HEAD` requests but will have
+    /// the response body removed. Make sure to add explicit `HEAD` routes
+    /// afterwards.
     pub fn get<H, B, T>(self, handler: H) -> OnMethod<IntoService<H, B, T>, Self>
     where
         H: Handler<B, T>,
     {
-        self.on(MethodFilter::GET, handler)
+        self.on(MethodFilter::GET | MethodFilter::HEAD, handler)
     }
 
     /// Chain an additional handler that will only accept `HEAD` requests.
@@ -624,6 +632,8 @@ where
     }
 
     fn call(&mut self, req: Request<B>) -> Self::Future {
+        let req_method = req.method().clone();
+
         let f = if self.method.matches(req.method()) {
             let fut = self.svc.clone().oneshot(req);
             RouteFuture::a(fut)
@@ -632,6 +642,9 @@ where
             RouteFuture::b(fut)
         };
 
-        future::OnMethodFuture { inner: f }
+        future::OnMethodFuture {
+            inner: f,
+            req_method,
+        }
     }
 }