Cookie Improvement

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"

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)
 

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();
+}

src/cookies.rs

@@ -1,46 +1,93 @@
 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<RwLock<CookieJar>>,
+}
+
+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<Cookie<'static>>;
+    fn get_cookie(&mut self, name: &str) -> Result<Option<Cookie<'static>>, 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<AppData> ExtractCookies for Context<AppData> {
-    fn cookie(&mut self, name: &str) -> Option<Cookie<'static>> {
+impl<AppData> CookiesExt for Context<AppData> {
+    fn get_cookie(&mut self, name: &str) -> Result<Option<Cookie<'static>>, 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::<CookieData>()
+            .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::<CookieData>()
+            .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::<CookieData>()
+            .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(())
     }
 }
 
 fn parse_from_header(s: &str) -> Result<CookieJar, ParseError> {
     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(())
     })?;
 

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<Data: Send + Sync + 'static> Middleware<Data> for CookiesMiddleware {
+    fn handle<'a>(
+        &'a self,
+        mut cx: Context<Data>,
+        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());
+    }
+
+}

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<AppData>: 'static + Send + Sync {