Automatically handle HEAD requests.

codegen/src/decorators/route.rs

@@ -285,3 +285,4 @@ method_decorator!(put_decorator, Put);
 method_decorator!(post_decorator, Post);
 method_decorator!(delete_decorator, Delete);
 method_decorator!(patch_decorator, Patch);
+method_decorator!(head_decorator, Head);

codegen/src/lib.rs

@@ -47,6 +47,7 @@ pub fn plugin_registrar(reg: &mut Registry) {
         "put" => put_decorator,
         "post" => post_decorator,
         "delete" => delete_decorator,
+        "head" => head_decorator,
         "patch" => patch_decorator
     );
 }

lib/src/response/response.rs

@@ -249,8 +249,8 @@ impl<'r> Response<'r> {
         // Looks crazy, right? Needed so Rust infers lifetime correctly. Weird.
         match self.body.as_mut() {
             Some(body) => Some(match body.as_mut() {
-                Body::Sized(b, u64) => Body::Sized(b, u64),
-                Body::Chunked(b, u64) => Body::Chunked(b, u64),
+                Body::Sized(b, size) => Body::Sized(b, size),
+                Body::Chunked(b, chunk_size) => Body::Chunked(b, chunk_size),
             }),
             None => None
         }
@@ -261,6 +261,19 @@ impl<'r> Response<'r> {
         self.body.take()
     }
 
+    // Removes any actual body, but leaves the size if it exists. Only meant to
+    // be used to handle HEAD requests automatically.
+    #[doc(hidden)]
+    #[inline(always)]
+    pub fn strip_body(&mut self) {
+        if let Some(body) = self.take_body() {
+            self.body = match body {
+                Body::Sized(_, n) => Some(Body::Sized(Box::new(io::empty()), n)),
+                Body::Chunked(..) => None
+            };
+        }
+    }
+
     #[inline(always)]
     pub fn set_sized_body<B>(&mut self, mut body: B)
         where B: io::Read + io::Seek + 'r

lib/src/rocket.rs

@@ -7,7 +7,7 @@ use std::io::{self, Write};
 use term_painter::Color::*;
 use term_painter::ToStyle;
 
-use logger;
+use {logger, handler};
 use ext::ReadExt;
 use config::{self, Config};
 use request::{Request, FormItems};
@@ -66,28 +66,8 @@ impl hyper::Handler for Rocket {
             }
         };
 
-        // Now that we've Rocket-ized everything, actually dispatch the request.
-        self.preprocess_request(&mut request, &data);
-        let mut response = match self.dispatch(&request, data) {
-            Ok(response) => response,
-            Err(status) => {
-                if status == Status::NotFound && request.method() == Method::Head {
-                    // FIXME: Handle unimplemented HEAD requests automatically.
-                    info_!("Redirecting to {}.", Green.paint(Method::Get));
-                }
-
-                let response = self.handle_error(status, &request);
-                return self.issue_response(response, res);
-            }
-        };
-
-        // We have a response from the user. Update the cookies in the header.
-        let cookie_delta = request.cookies().delta();
-        if cookie_delta.len() > 0 {
-            response.adjoin_header(header::SetCookie(cookie_delta));
-        }
-
-        // Actually write out the response.
+        // Dispatch the request to get a response, then write that response out.
+        let response = self.dispatch(&mut request, data);
         return self.issue_response(response, res);
     }
 }
@@ -175,14 +155,53 @@ impl Rocket {
         }
     }
 
+    #[doc(hidden)]
+    pub fn dispatch<'r>(&self, request: &'r mut Request, data: Data) -> Response<'r> {
+        // Do a bit of preprocessing before routing.
+        self.preprocess_request(request, &data);
+
+        // Route the request to get a response.
+        match self.route(request, data) {
+            Outcome::Success(mut response) => {
+                let cookie_delta = request.cookies().delta();
+                if cookie_delta.len() > 0 {
+                    response.adjoin_header(header::SetCookie(cookie_delta));
+                }
+
+                response
+            }
+            Outcome::Forward(data) => {
+                // Rust thinks `request` is still borrowed here, but it's
+                // obviously not (data has nothing to do with it), so we
+                // convince it to give us another mutable reference.
+                // TODO: Pay the cost to copy Request into UnsafeCell?
+                let request: &'r mut Request = unsafe {
+                    &mut *(request as *const Request as *mut Request)
+                };
+
+                if request.method() == Method::Head {
+                    info_!("Autohandling {} request.", White.paint("HEAD"));
+                    request.set_method(Method::Get);
+                    let mut response = self.dispatch(request, data);
+                    response.strip_body();
+                    response
+                } else {
+                    self.handle_error(Status::NotFound, request)
+                }
+            }
+            Outcome::Failure(status) => self.handle_error(status, request),
+        }
+    }
+
     /// Tries to find a `Responder` for a given `request`. It does this by
     /// routing the request and calling the handler for each matching route
     /// until one of the handlers returns success or failure. If a handler
     /// returns a failure, or there are no matching handlers willing to accept
     /// the request, this function returns an `Err` with the status code.
     #[doc(hidden)]
-    pub fn dispatch<'r>(&self, request: &'r Request, mut data: Data)
-            -> Result<Response<'r>, Status> {
+    #[inline]
+    pub fn route<'r>(&self, request: &'r Request, mut data: Data)
+            -> handler::Outcome<'r> {
         // Go through the list of matching routes until we fail or succeed.
         info!("{}:", request);
         let matches = self.router.route(&request);
@@ -198,14 +217,14 @@ impl Rocket {
             // to be forwarded. If it does, continue the loop to try again.
             info_!("{} {}", White.paint("Outcome:"), outcome);
             match outcome {
-                Outcome::Success(response) => return Ok(response),
-                Outcome::Failure(status) => return Err(status),
+                o@Outcome::Success(_) => return o,
+                o@Outcome::Failure(_) => return o,
                 Outcome::Forward(unused_data) => data = unused_data,
             };
         }
 
         error_!("No matching routes for {}.", request);
-        Err(Status::NotFound)
+        Outcome::Forward(data)
     }
 
     // TODO: DOC.

lib/src/testing.rs

@@ -222,9 +222,6 @@ impl MockRequest {
     /// ```
     pub fn dispatch_with<'r>(&'r mut self, rocket: &Rocket) -> Response<'r> {
         let data = ::std::mem::replace(&mut self.data, Data::new(vec![]));
-        match rocket.dispatch(&self.request, data) {
-            Ok(response) => response,
-            Err(status) => rocket.handle_error(status, &self.request)
-        }
+        rocket.dispatch(&mut self.request, data)
     }
 }

lib/tests/head_handling.rs

@@ -0,0 +1,72 @@
+#![feature(plugin)]
+#![plugin(rocket_codegen)]
+
+extern crate rocket;
+
+use rocket::Route;
+use rocket::response::{status, content};
+use rocket::http::ContentType;
+
+#[get("/empty")]
+fn empty() -> status::NoContent {
+    status::NoContent
+}
+
+#[get("/")]
+fn index() -> &'static str {
+    "Hello, world!"
+}
+
+#[head("/other")]
+fn other() -> content::JSON<()> {
+    content::JSON(())
+}
+
+fn routes() -> Vec<Route> {
+    routes![index, empty, other]
+}
+
+use rocket::testing::MockRequest;
+use rocket::http::Method::*;
+use rocket::http::Status;
+use rocket::response::Body;
+
+#[test]
+fn auto_head() {
+    let rocket = rocket::ignite().mount("/", routes());
+
+    let mut req = MockRequest::new(Head, "/");
+    let mut response = req.dispatch_with(&rocket);
+
+    assert_eq!(response.status(), Status::Ok);
+    if let Some(body) = response.body() {
+        match body {
+            Body::Sized(_, n) => assert_eq!(n, "Hello, world!".len() as u64),
+            _ => panic!("Expected a sized body!")
+        }
+
+        assert_eq!(body.to_string(), Some("".to_string()));
+    } else {
+        panic!("Expected an empty body!")
+    }
+
+
+    let content_type: Vec<_> = response.get_header_values("Content-Type").collect();
+    assert_eq!(content_type, vec![ContentType::Plain.to_string()]);
+
+    let mut req = MockRequest::new(Head, "/empty");
+    let response = req.dispatch_with(&rocket);
+    assert_eq!(response.status(), Status::NoContent);
+}
+
+#[test]
+fn user_head() {
+    let rocket = rocket::ignite().mount("/", routes());
+    let mut req = MockRequest::new(Head, "/other");
+    let response = req.dispatch_with(&rocket);
+
+    assert_eq!(response.status(), Status::Ok);
+
+    let content_type: Vec<_> = response.get_header_values("Content-Type").collect();
+    assert_eq!(content_type, vec![ContentType::JSON.to_string()]);
+}