From 9629b402024709b50ef2ac10e0be00c042ab73ed Mon Sep 17 00:00:00 2001 From: Sergio Benitez Date: Sat, 14 Apr 2018 19:53:36 -0700 Subject: [PATCH] Clear flash cookies only after they're inspected. Resolves #512. Resolves #466. --- examples/session/templates/login.html.hbs | 7 ++- lib/src/request/mod.rs | 5 +- lib/src/response/flash.rs | 74 +++++++++++++++-------- lib/src/response/mod.rs | 3 +- lib/tests/flash-lazy-removes-issue-466.rs | 63 +++++++++++++++++++ 5 files changed, 119 insertions(+), 33 deletions(-) create mode 100644 lib/tests/flash-lazy-removes-issue-466.rs diff --git a/examples/session/templates/login.html.hbs b/examples/session/templates/login.html.hbs index 504a2bcc..afe6276e 100644 --- a/examples/session/templates/login.html.hbs +++ b/examples/session/templates/login.html.hbs @@ -7,10 +7,11 @@

Rocket Session: Please Login

+ +

Please login to continue.

+ {{#if flash}} -

{{ flash }}

- {{else}} -

Please login to continue.

+

Error: {{ flash }}

{{/if}}
diff --git a/lib/src/request/mod.rs b/lib/src/request/mod.rs index 723e07eb..b4cc2be5 100644 --- a/lib/src/request/mod.rs +++ b/lib/src/request/mod.rs @@ -15,6 +15,5 @@ pub use self::param::{FromParam, FromSegments}; pub use self::form::{Form, LenientForm, FromForm, FromFormValue, FormItems}; pub use self::state::State; -/// Type alias to retrieve [Flash](/rocket/response/struct.Flash.html) messages -/// from a request. -pub type FlashMessage = ::response::Flash<()>; +#[doc(inline)] +pub use response::flash::FlashMessage; diff --git a/lib/src/response/flash.rs b/lib/src/response/flash.rs index 17d37da7..5b0bc356 100644 --- a/lib/src/response/flash.rs +++ b/lib/src/response/flash.rs @@ -6,6 +6,7 @@ use outcome::IntoOutcome; use response::{Response, Responder}; use request::{self, Request, FromRequest}; use http::{Status, Cookie}; +use std::sync::atomic::{AtomicBool, Ordering}; // The name of the actual flash cookie. const FLASH_COOKIE_NAME: &'static str = "_flash"; @@ -87,9 +88,26 @@ const FLASH_COOKIE_NAME: &'static str = "_flash"; pub struct Flash { name: String, message: String, - responder: R, + consumed: AtomicBool, + inner: R, } +/// Type alias to retrieve [`Flash`](/rocket/response/struct.Flash.html) +/// messages from a request. +/// +/// # Flash Cookie +/// +/// A `FlashMessage` holds the parsed contents of the flash cookie. As long as +/// there is a flash cookie present (set by the `Flash` `Responder`), a +/// `FlashMessage` reuqest guard will succeed. +/// +/// The flash cookie is cleared if either the [`name()`] or [`msg()`] method is +/// called. If neither method is called, the flash cookie is not cleared. +/// +/// [`name()`]: /rocket/response.struct.Flash.html#method.name +/// [`msg()`]: /rocket/response.struct.Flash.html#method.msg +pub type FlashMessage<'a, 'r> = ::response::Flash<&'a Request<'r>>; + impl<'r, R: Responder<'r>> Flash { /// Constructs a new `Flash` message with the given `name`, `msg`, and /// underlying `responder`. @@ -109,7 +127,8 @@ impl<'r, R: Responder<'r>> Flash { Flash { name: name.as_ref().to_string(), message: msg.as_ref().to_string(), - responder: res, + consumed: AtomicBool::default(), + inner: res, } } @@ -183,45 +202,56 @@ impl<'r, R: Responder<'r>> Flash { impl<'r, R: Responder<'r>> Responder<'r> for Flash { fn respond_to(self, req: &Request) -> Result, Status> { trace_!("Flash: setting message: {}:{}", self.name, self.message); - let cookie = self.cookie(); - Response::build_from(self.responder.respond_to(req)?) - .header_adjoin(&cookie) - .ok() + req.cookies().add(self.cookie()); + self.inner.respond_to(req) } } -impl Flash<()> { - /// Constructs a new message with the given name and message. - fn named(name: &str, msg: &str) -> Flash<()> { +impl<'a, 'r> Flash<&'a Request<'r>> { + /// Constructs a new message with the given name and message for the given + /// request. + fn named(name: &str, msg: &str, request: &'a Request<'r>) -> Flash<&'a Request<'r>> { Flash { name: name.to_string(), message: msg.to_string(), - responder: (), + consumed: AtomicBool::new(false), + inner: request, + } + } + + // Clears the request cookie if it hasn't already been cleared. + fn clear_cookie_if_needed(&self) { + // Remove the cookie if it hasn't already been removed. + if !self.consumed.swap(true, Ordering::Relaxed) { + let cookie = Cookie::build(FLASH_COOKIE_NAME, "").path("/").finish(); + self.inner.cookies().remove(cookie); } } /// Returns the `name` of this message. pub fn name(&self) -> &str { - self.name.as_str() + self.clear_cookie_if_needed(); + &self.name } /// Returns the `msg` contents of this message. pub fn msg(&self) -> &str { - self.message.as_str() + self.clear_cookie_if_needed(); + &self.message } } -/// Retrieves a flash message from a flash cookie and deletes the flash cookie. -/// If there is no flash cookie, an empty `Err` is returned. +/// Retrieves a flash message from a flash cookie. If there is no flash cookie, +/// or if the flash cookie is malformed, an empty `Err` is returned. /// /// The suggested use is through an `Option` and the `FlashMessage` type alias /// in `request`: `Option`. -impl<'a, 'r> FromRequest<'a, 'r> for Flash<()> { +impl<'a, 'r> FromRequest<'a, 'r> for Flash<&'a Request<'r>> { type Error = (); fn from_request(request: &'a Request<'r>) -> request::Outcome { trace_!("Flash: attemping to retrieve message."); - let r = request.cookies().get(FLASH_COOKIE_NAME).ok_or(()).and_then(|cookie| { + request.cookies().get(FLASH_COOKIE_NAME).ok_or(()).and_then(|cookie| { trace_!("Flash: retrieving message: {:?}", cookie); // Parse the flash message. @@ -233,15 +263,7 @@ impl<'a, 'r> FromRequest<'a, 'r> for Flash<()> { let name_len: usize = len_str.parse().map_err(|_| ())?; let (name, msg) = (&rest[..name_len], &rest[name_len..]); - Ok(Flash::named(name, msg)) - }); - - // If we found a flash cookie, delete it from the jar. - if r.is_ok() { - let cookie = Cookie::build(FLASH_COOKIE_NAME, "").path("/").finish(); - request.cookies().remove(cookie); - } - - r.into_outcome(Status::BadRequest) + Ok(Flash::named(name, msg, request)) + }).into_outcome(Status::BadRequest) } } diff --git a/lib/src/response/mod.rs b/lib/src/response/mod.rs index 486f5299..638d5a4d 100644 --- a/lib/src/response/mod.rs +++ b/lib/src/response/mod.rs @@ -21,12 +21,13 @@ mod responder; mod redirect; -mod flash; mod named_file; mod stream; mod response; mod failure; +pub(crate) mod flash; + pub mod content; pub mod status; diff --git a/lib/tests/flash-lazy-removes-issue-466.rs b/lib/tests/flash-lazy-removes-issue-466.rs new file mode 100644 index 00000000..7d6703c1 --- /dev/null +++ b/lib/tests/flash-lazy-removes-issue-466.rs @@ -0,0 +1,63 @@ +#![feature(plugin, decl_macro, custom_derive)] +#![plugin(rocket_codegen)] + +extern crate rocket; + +use rocket::request::FlashMessage; +use rocket::response::Flash; + +const FLASH_MESSAGE: &str = "Hey! I'm a flash message. :)"; + +#[post("/")] +fn set() -> Flash<&'static str> { + Flash::success("This is the page.", FLASH_MESSAGE) +} + +#[get("/unused")] +fn unused(flash: Option) -> Option<()> { + flash.map(|_| ()) +} + +#[get("/use")] +fn used(flash: Option) -> Option { + flash.map(|flash| flash.msg().into()) +} + +mod flash_lazy_remove_tests { + use rocket::local::Client; + use rocket::http::Status; + + #[test] + fn test() { + use super::*; + let r = rocket::ignite().mount("/", routes![set, unused, used]); + let client = Client::new(r).unwrap(); + + // Ensure the cookie's not there at first. + let response = client.get("/unused").dispatch(); + assert_eq!(response.status(), Status::NotFound); + + // Set the flash cookie. + client.post("/").dispatch(); + + // Try once. + let response = client.get("/unused").dispatch(); + assert_eq!(response.status(), Status::Ok); + + // Try again; should still be there. + let response = client.get("/unused").dispatch(); + assert_eq!(response.status(), Status::Ok); + + // Now use it. + let mut response = client.get("/use").dispatch(); + assert_eq!(response.body_string(), Some(FLASH_MESSAGE.into())); + + // Now it should be gone. + let response = client.get("/unused").dispatch(); + assert_eq!(response.status(), Status::NotFound); + + // Still gone. + let response = client.get("/use").dispatch(); + assert_eq!(response.status(), Status::NotFound); + } +}