From 2e25ce04dc49b5c1a0ee5ccff70b577ad795e16d Mon Sep 17 00:00:00 2001 From: Sergio Benitez Date: Fri, 16 Dec 2016 05:17:16 -0800 Subject: [PATCH] Automatically handle HEAD requests. --- codegen/src/decorators/route.rs | 1 + codegen/src/lib.rs | 1 + lib/src/response/response.rs | 17 +++++++- lib/src/rocket.rs | 75 +++++++++++++++++++++------------ lib/src/testing.rs | 5 +-- lib/tests/head_handling.rs | 72 +++++++++++++++++++++++++++++++ 6 files changed, 137 insertions(+), 34 deletions(-) create mode 100644 lib/tests/head_handling.rs diff --git a/codegen/src/decorators/route.rs b/codegen/src/decorators/route.rs index 349e3bed..814ff4e1 100644 --- a/codegen/src/decorators/route.rs +++ b/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); diff --git a/codegen/src/lib.rs b/codegen/src/lib.rs index e5caf519..5b5c7e74 100644 --- a/codegen/src/lib.rs +++ b/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 ); } diff --git a/lib/src/response/response.rs b/lib/src/response/response.rs index a2e2fd44..275384ca 100644 --- a/lib/src/response/response.rs +++ b/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(&mut self, mut body: B) where B: io::Read + io::Seek + 'r diff --git a/lib/src/rocket.rs b/lib/src/rocket.rs index 8c7b354d..de4127f9 100644 --- a/lib/src/rocket.rs +++ b/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, 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. diff --git a/lib/src/testing.rs b/lib/src/testing.rs index 9731f576..75182fca 100644 --- a/lib/src/testing.rs +++ b/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) } } diff --git a/lib/tests/head_handling.rs b/lib/tests/head_handling.rs new file mode 100644 index 00000000..f92c660a --- /dev/null +++ b/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 { + 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()]); +}