Cookie Improvement

Cargo.toml

@@ -13,7 +13,7 @@ repository = "https://github.com/rustasync/tide"
 version = "0.0.5"
 
 [dependencies]
-cookie = "0.11"
+cookie = { version="0.11", features = ["percent-encode"] }
 futures-preview = "0.3.0-alpha.13"
 fnv = "1.0.6"
 http = "0.1"

examples/cookies.rs

@@ -0,0 +1,28 @@
+#![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<()>) -> String {
+    cx.set_cookie(Cookie::new("hello", "world")).unwrap();
+    format!("hello cookies: {:?}", cx.get_cookie("hello"))
+}
+async fn remove_cookie(mut cx: Context<()>) -> String {
+    cx.remove_cookie(Cookie::new("hello", "world")).unwrap();
+    format!("hello cookies: {:?}", cx.get_cookie("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,88 @@
 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>;
+    fn set_cookie(&mut self, cookie: Cookie<'static>) -> Result<(), StringError>;
+    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()))?;
+
+        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())
+    }
 
-        cookie
+    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,52 @@
+use crate::cookies::CookieData;
+use futures::future::FutureObj;
+use http::header::HeaderValue;
+
+use crate::{
+    middleware::{Middleware, Next},
+    Context, Response,
+};
+
+#[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 Log 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
+        }
+    }
+}

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 {