diff --git a/Cargo.toml b/Cargo.toml index 5b658dca7..9613d0f30 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,7 @@ repository = "https://github.com/rustasync/tide" version = "0.1.1" [dependencies] -cookie = "0.11" +cookie = { version="0.11", features = ["percent-encode"] } futures-preview = "0.3.0-alpha.14" fnv = "1.0.6" http = "0.1" diff --git a/README.md b/README.md index 969a16a2a..47078dc73 100644 --- a/README.md +++ b/README.md @@ -72,7 +72,7 @@ __More Examples__ - [Body Types](https://github.com/rustasync/tide/blob/master/examples/body_types.rs) - [Multipart Form](https://github.com/rustasync/tide/tree/master/examples/multipart-form/main.rs) - [Catch All](https://github.com/rustasync/tide/tree/master/examples/catch_all.rs) -- [Cookie Extractor](https://github.com/rustasync/tide/tree/master/examples/cookie_extractor.rs) +- [Cookies](https://github.com/rustasync/tide/tree/master/examples/cookies.rs) - [Default Headers](https://github.com/rustasync/tide/tree/master/examples/default_headers.rs) - [GraphQL](https://github.com/rustasync/tide/tree/master/examples/graphql.rs) diff --git a/examples/cookie_extractor.rs b/examples/cookie_extractor.rs deleted file mode 100644 index eb31b42cd..000000000 --- a/examples/cookie_extractor.rs +++ /dev/null @@ -1,15 +0,0 @@ -#![feature(async_await, futures_api)] - -use tide::{cookies::ExtractCookies, Context}; - -/// Tide will use the the `Cookies`'s `Extract` implementation to build this parameter. -/// -async fn hello_cookies(mut cx: Context<()>) -> String { - format!("hello cookies: {:?}", cx.cookie("hello")) -} - -fn main() { - let mut app = tide::App::new(()); - app.at("/").get(hello_cookies); - app.serve("127.0.0.1:8000").unwrap(); -} diff --git a/examples/cookies.rs b/examples/cookies.rs new file mode 100644 index 000000000..7f9041bb0 --- /dev/null +++ b/examples/cookies.rs @@ -0,0 +1,26 @@ +#![feature(async_await, futures_api)] + +use cookie::Cookie; +use tide::{cookies::CookiesExt, middleware::CookiesMiddleware, Context}; + +/// Tide will use the the `Cookies`'s `Extract` implementation to build this parameter. +/// +async fn retrieve_cookie(mut cx: Context<()>) -> String { + format!("hello cookies: {:?}", cx.get_cookie("hello").unwrap()) +} +async fn set_cookie(mut cx: Context<()>) { + cx.set_cookie(Cookie::new("hello", "world")).unwrap(); +} +async fn remove_cookie(mut cx: Context<()>) { + cx.remove_cookie(Cookie::named("hello")).unwrap(); +} + +fn main() { + let mut app = tide::App::new(()); + app.middleware(CookiesMiddleware::new()); + + app.at("/").get(retrieve_cookie); + app.at("/set").get(set_cookie); + app.at("/remove").get(remove_cookie); + app.serve("127.0.0.1:8000").unwrap(); +} diff --git a/src/cookies.rs b/src/cookies.rs index c5f67ee6b..06af7de94 100644 --- a/src/cookies.rs +++ b/src/cookies.rs @@ -1,37 +1,85 @@ use cookie::{Cookie, CookieJar, ParseError}; +use crate::error::StringError; use crate::Context; +use http::HeaderMap; +use std::sync::{Arc, RwLock}; + +const MIDDLEWARE_MISSING_MSG: &str = + "CookiesMiddleware must be used to populate request and response cookies"; /// A representation of cookies which wraps `CookieJar` from `cookie` crate /// -/// Currently this only exposes getting cookie by name but future enhancements might allow more -/// operations -struct CookieData { - content: CookieJar, +#[derive(Debug)] +pub(crate) struct CookieData { + pub(crate) content: Arc>, +} + +impl CookieData { + pub fn from_headers(headers: &HeaderMap) -> Self { + CookieData { + content: Arc::new(RwLock::new( + headers + .get(http::header::COOKIE) + .and_then(|raw| parse_from_header(raw.to_str().unwrap()).ok()) + .unwrap_or_default(), + )), + } + } } /// An extension to `Context` that provides cached access to cookies -pub trait ExtractCookies { +pub trait CookiesExt { /// returns a `Cookie` by name of the cookie - fn cookie(&mut self, name: &str) -> Option>; + fn get_cookie(&mut self, name: &str) -> Result>, StringError>; + + /// Add cookie to the cookie jar + fn set_cookie(&mut self, cookie: Cookie<'static>) -> Result<(), StringError>; + + /// Removes the cookie. This instructs the `CookiesMiddleware` to send a cookie with empty value + /// in the response. + fn remove_cookie(&mut self, cookie: Cookie<'static>) -> Result<(), StringError>; } -impl ExtractCookies for Context { - fn cookie(&mut self, name: &str) -> Option> { +impl CookiesExt for Context { + fn get_cookie(&mut self, name: &str) -> Result>, StringError> { let cookie_data = self - .extensions_mut() - .remove() - .unwrap_or_else(|| CookieData { - content: self - .headers() - .get("tide-cookie") - .and_then(|raw| parse_from_header(raw.to_str().unwrap()).ok()) - .unwrap_or_default(), - }); - let cookie = cookie_data.content.get(name).cloned(); - self.extensions_mut().insert(cookie_data); + .extensions() + .get::() + .ok_or_else(|| StringError(MIDDLEWARE_MISSING_MSG.to_owned()))?; - cookie + let arc_jar = cookie_data.content.clone(); + let locked_jar = arc_jar + .read() + .map_err(|e| StringError(format!("Failed to get write lock: {}", e)))?; + Ok(locked_jar.get(name).cloned()) + } + + fn set_cookie(&mut self, cookie: Cookie<'static>) -> Result<(), StringError> { + let cookie_data = self + .extensions() + .get::() + .ok_or_else(|| StringError(MIDDLEWARE_MISSING_MSG.to_owned()))?; + let jar = cookie_data.content.clone(); + let mut locked_jar = jar + .write() + .map_err(|e| StringError(format!("Failed to get write lock: {}", e)))?; + locked_jar.add(cookie); + Ok(()) + } + + fn remove_cookie(&mut self, cookie: Cookie<'static>) -> Result<(), StringError> { + let cookie_data = self + .extensions() + .get::() + .ok_or_else(|| StringError(MIDDLEWARE_MISSING_MSG.to_owned()))?; + + let jar = cookie_data.content.clone(); + let mut locked_jar = jar + .write() + .map_err(|e| StringError(format!("Failed to get write lock: {}", e)))?; + locked_jar.remove(cookie); + Ok(()) } } @@ -39,8 +87,7 @@ fn parse_from_header(s: &str) -> Result { let mut jar = CookieJar::new(); s.split(';').try_for_each(|s| -> Result<_, ParseError> { - jar.add(Cookie::parse(s.trim().to_owned())?); - + jar.add_original(Cookie::parse(s.trim().to_owned())?); Ok(()) })?; diff --git a/src/middleware/cookies.rs b/src/middleware/cookies.rs new file mode 100644 index 000000000..ea0e97b2e --- /dev/null +++ b/src/middleware/cookies.rs @@ -0,0 +1,170 @@ +use crate::cookies::CookieData; +use futures::future::FutureObj; +use http::header::HeaderValue; + +use crate::{ + middleware::{Middleware, Next}, + Context, Response, +}; + +/// Middleware to work with cookies. +/// +/// [`CookiesMiddleware`] along with [`CookiesExt`](crate::cookies::CookiesExt) provide smooth +/// access to request cookies and setting/removing cookies from response. This leverages the +/// [cookie](https://crates.io/crates/cookie) crate. +/// This middleware parses cookies from request and caches them in the extension. Once the request +/// is processed by endpoints and other middlewares, all the added and removed cookies are set on +/// on the response. You will need to add this middle before any other middlewares that might need +/// to access Cookies. +#[derive(Clone, Default)] +pub struct CookiesMiddleware {} + +impl CookiesMiddleware { + pub fn new() -> Self { + Self {} + } +} + +impl Middleware for CookiesMiddleware { + fn handle<'a>( + &'a self, + mut cx: Context, + next: Next<'a, Data>, + ) -> FutureObj<'a, Response> { + box_async! { + let cookie_data = cx + .extensions_mut() + .remove() + .unwrap_or_else(|| CookieData::from_headers(cx.headers())); + + let cookie_jar = cookie_data.content.clone(); + + cx.extensions_mut().insert(cookie_data); + let mut res = await!(next.run(cx)); + let headers = res.headers_mut(); + for cookie in cookie_jar.read().unwrap().delta() { + let hv = HeaderValue::from_str(cookie.encoded().to_string().as_str()); + if let Ok(val) = hv { + headers.append(http::header::SET_COOKIE, val); + } else { + // TODO It would be useful to log this error here. + return http::Response::builder() + .status(http::status::StatusCode::INTERNAL_SERVER_ERROR) + .header("Content-Type", "text/plain; charset=utf-8") + .body(http_service::Body::empty()) + .unwrap(); + } + } + res + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{cookies::CookiesExt, Context}; + use cookie::Cookie; + use futures::executor::block_on; + use http_service::Body; + use http_service_mock::make_server; + + static COOKIE_NAME: &str = "testCookie"; + + /// Tide will use the the `Cookies`'s `Extract` implementation to build this parameter. + async fn retrieve_cookie(mut cx: Context<()>) -> String { + format!("{}", cx.get_cookie(COOKIE_NAME).unwrap().unwrap().value()) + } + async fn set_cookie(mut cx: Context<()>) { + cx.set_cookie(Cookie::new(COOKIE_NAME, "NewCookieValue")) + .unwrap(); + } + async fn remove_cookie(mut cx: Context<()>) { + cx.remove_cookie(Cookie::named(COOKIE_NAME)).unwrap(); + } + + async fn set_multiple_cookie(mut cx: Context<()>) { + cx.set_cookie(Cookie::new("C1", "V1")).unwrap(); + cx.set_cookie(Cookie::new("C2", "V2")).unwrap(); + } + + fn app() -> crate::App<()> { + let mut app = crate::App::new(()); + app.middleware(CookiesMiddleware::new()); + + app.at("/get").get(retrieve_cookie); + app.at("/set").get(set_cookie); + app.at("/remove").get(remove_cookie); + app.at("/multi").get(set_multiple_cookie); + app + } + + fn make_request(endpoint: &str) -> Response { + let app = app(); + let mut server = make_server(app.into_http_service()).unwrap(); + let req = http::Request::get(endpoint) + .header(http::header::COOKIE, "testCookie=RequestCookieValue") + .body(Body::empty()) + .unwrap(); + let res = server.simulate(req).unwrap(); + res + } + + #[test] + fn successfully_retrieve_request_cookie() { + let res = make_request("/get"); + assert_eq!(res.status(), 200); + let body = block_on(res.into_body().into_vec()).unwrap(); + assert_eq!(&*body, &*b"RequestCookieValue"); + } + + #[test] + fn successfully_set_cookie() { + let res = make_request("/set"); + assert_eq!(res.status(), 200); + let test_cookie_header = res.headers().get(http::header::SET_COOKIE).unwrap(); + assert_eq!( + test_cookie_header.to_str().unwrap(), + "testCookie=NewCookieValue" + ); + } + + #[test] + fn successfully_remove_cookie() { + let res = make_request("/remove"); + assert_eq!(res.status(), 200); + let test_cookie_header = res.headers().get(http::header::SET_COOKIE).unwrap(); + assert!(test_cookie_header + .to_str() + .unwrap() + .starts_with("testCookie=;")); + let cookie = Cookie::parse_encoded(test_cookie_header.to_str().unwrap()).unwrap(); + assert_eq!(cookie.name(), COOKIE_NAME); + assert_eq!(cookie.value(), ""); + assert_eq!(cookie.http_only(), None); + assert_eq!(cookie.max_age().unwrap().num_nanoseconds(), Some(0)); + } + + #[test] + fn successfully_set_multiple_cookies() { + let res = make_request("/multi"); + assert_eq!(res.status(), 200); + let cookie_header = res.headers().get_all(http::header::SET_COOKIE); + let mut iter = cookie_header.iter(); + + let cookie1 = iter.next().unwrap(); + let cookie2 = iter.next().unwrap(); + + //Headers can be out of order + if cookie1.to_str().unwrap().starts_with("C1") { + assert_eq!(cookie1, "C1=V1"); + assert_eq!(cookie2, "C2=V2"); + } else { + assert_eq!(cookie2, "C1=V1"); + assert_eq!(cookie1, "C2=V2"); + } + + assert!(iter.next().is_none()); + } + +} diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs index 23f8823f0..d48c1e16f 100644 --- a/src/middleware/mod.rs +++ b/src/middleware/mod.rs @@ -3,10 +3,11 @@ use std::sync::Arc; use crate::{endpoint::DynEndpoint, Context, Response}; +mod cookies; mod default_headers; mod logger; -pub use self::{default_headers::DefaultHeaders, logger::RootLogger}; +pub use self::{cookies::CookiesMiddleware, default_headers::DefaultHeaders, logger::RootLogger}; /// Middleware that wraps around remaining middleware chain. pub trait Middleware: 'static + Send + Sync {