obsidian-rs · plwai · Feb 11, 2020 · Feb 11, 2020 · May 1, 2020 · May 13, 2020
diff --git a/Cargo.toml b/Cargo.toml
@@ -18,6 +18,14 @@ path = 'examples/json.rs'
 name = 'app_state'
 path = 'examples/app_state.rs'
 
+[[example]]
+name = 'cookies'
+path = 'examples/cookies.rs'
+
+[[example]]
+name = 'session'
+path = 'examples/session.rs'
+
 [package]
 name = 'obsidian'
 version = '0.3.0-dev'
@@ -48,6 +56,10 @@ url = '2.1.1'
 async-std = "1.5.0"
 async-trait = "0.1.30"
 colored = "1.9.3"
+cookie = "0.13.3"
+chrono = "0.4.11"
+lazy_static = "1.4.0"
+rand = "0.7.3"
 
 [dependencies.serde]
 version = '1.0.106'

diff --git a/examples/cookies.rs b/examples/cookies.rs
@@ -0,0 +1,35 @@
+use obsidian::{context::Context, App};
+use obsidian::middleware::cookie_parser::CookieParser;
+use cookie::Cookie;
+
+#[tokio::main]
+async fn main() {
+    let mut app: App = App::new();
+    let addr = ([127, 0, 0, 1], 3000).into();
+
+    app.get("/", |ctx: Context| async { 
+        if let Some(cookies) = ctx.cookie("cookie_name") {
+            let cookies2 = ctx.cookie("cookie_name2").expect("cookie_name2 should be set");
+            let response = format!("{}={}; {}={};", cookies.name(), cookies.value(), cookies2.name(), cookies2.value());
+
+            ctx
+            .build(response)
+            .ok() 
+        }
+        else {
+            ctx
+            .build("Set cookies!")
+            .with_cookie(Cookie::new("cookie_name", "cookie_value"))
+            .with_cookie(Cookie::new("cookie_name2", "cookie_value2"))
+            .ok() 
+        }
+    });
+
+    let cookie_parser = CookieParser::new();
+    app.use_service(cookie_parser);
+
+    app.listen(&addr, || {
+        println!("server is listening to {}", &addr);
+    })
+    .await;
+}
diff --git a/examples/main.rs b/examples/main.rs
@@ -132,8 +132,8 @@ ctx.build(Response::ok().html("<!DOCTYPE html><html><head><link rel=\"shotcut ic
 
         ctx.build(
             Response::created()
-                .with_headers(standard_headers)
-                .with_headers_str(custom_headers)
+                .with_headers(&standard_headers)
+                .with_headers_str(&custom_headers)
                 .json(point),
         )
         .ok()
@@ -152,8 +152,8 @@ ctx.build(Response::ok().html("<!DOCTYPE html><html><head><link rel=\"shotcut ic
         ];
 
         ctx.build("Hello World")
-            .with_headers(standard_headers)
-            .with_headers_str(custom_headers)
+            .with_headers(&standard_headers)
+            .with_headers_str(&custom_headers)
             .ok()
     });
 

diff --git a/examples/session.rs b/examples/session.rs
@@ -0,0 +1,32 @@
+use obsidian::{context::Context, App};
+use obsidian::context::memory_session::MemorySessionStorage;
+use obsidian::middleware::{cookie_parser::CookieParser, cookie_session::CookieSession};
+
+#[tokio::main]
+async fn main() {
+    let mut app: App = App::new();
+    let addr = ([127, 0, 0, 1], 3000).into();
+
+    app.get("/", |mut ctx: Context| async { 
+        let session_content = match ctx.session() {
+            Some(session) => {
+                session.get("test").unwrap()
+            },
+            _ => "",
+        }.to_owned();
+
+        ctx.session_set("test", "session is set!");
+        ctx.build(session_content).ok()
+    });
+
+    let cookie_parser = CookieParser::new();
+    let cookie_session = CookieSession::new(MemorySessionStorage::new());
+
+    app.use_service(cookie_parser);
+    app.use_service(cookie_session);
+
+    app.listen(&addr, || {
+        println!("server is listening to {}", &addr);
+    })
+    .await;
+}
diff --git a/src/app.rs b/src/app.rs
@@ -215,17 +215,39 @@ impl AppServer {
                 let route_result = executor.next(context).await;
 
                 let route_response = match route_result {
-                    Ok(ctx) => {
+                    Ok(mut ctx) => {
                         let mut res = Response::builder();
                         if let Some(response) = ctx.take_response() {
                             if let Some(headers) = response.headers() {
                                 if let Some(response_headers) = res.headers_mut() {
                                     headers.iter().for_each(move |(key, value)| {
-                                        response_headers
-                                            .insert(key, header::HeaderValue::from_static(value));
+                                        if let Ok(header_value) =
+                                            header::HeaderValue::from_str(value)
+                                        {
+                                            response_headers.insert(key, header_value);
+                                        }
                                     });
                                 }
                             }
+
+                            if let Some(cookies) = response.cookies() {
+                                if let Some(response_headers) = res.headers_mut() {
+                                    cookies.iter().for_each(move |cookie| {
+                                        if let Ok(header_value) =
+                                            header::HeaderValue::from_str(&cookie.to_string())
+                                        {
+                                            if response_headers.contains_key(header::SET_COOKIE) {
+                                                response_headers
+                                                    .append(header::SET_COOKIE, header_value);
+                                            } else {
+                                                response_headers
+                                                    .insert(header::SET_COOKIE, header_value);
+                                            }
+                                        }
+                                    });
+                                }
+                            }
+
                             res.status(response.status()).body(response.body())
                         } else {
                             // No response found

diff --git a/src/context.rs b/src/context.rs
@@ -1,3 +1,8 @@
+pub mod cookies;
+pub mod memory_session;
+pub mod session;
+
+use cookie::Cookie;
 use http::Extensions;
 use hyper::{body, body::Buf};
 use serde::de::DeserializeOwned;
@@ -9,6 +14,8 @@ use std::collections::HashMap;
 use std::convert::From;
 use std::str::FromStr;
 
+use self::cookies::CookieParserData;
+use self::session::SessionData;
 use crate::router::{from_cow_map, ContextResult, Responder, Response};
 use crate::ObsidianError;
 use crate::{
@@ -256,14 +263,43 @@ impl Context {
         unimplemented!()
     }
 
+    pub fn cookie(&self, name: &str) -> Option<&Cookie> {
+        if let Some(cookie_data) = self.get::<CookieParserData>() {
+            return cookie_data.cookie_jar().get(name);
+        }
+
+        None
+    }
+
+    pub fn session(&self) -> Option<&SessionData> {
+        self.get::<SessionData>()
+    }
+
+    pub fn session_mut(&mut self) -> Option<&mut SessionData> {
+        self.get_mut::<SessionData>()
+    }
+
+    pub fn session_set(&mut self, name: &str, value: &str) {
+        match self.get_mut::<SessionData>() {
+            Some(session) => {
+                session.set(name, value);
+            }
+            _ => {
+                let mut session = SessionData::new();
+                session.set(name, value);
+                self.add(session);
+            }
+        }
+    }
+
     /// Consumes body of the request and replace it with empty body.
     pub fn take_body(&mut self) -> Body {
         std::mem::replace(self.request.body_mut(), Body::empty())
     }
 
-    /// Take response
-    pub fn take_response(self) -> Option<Response> {
-        self.response
+    /// Consumes response
+    pub fn take_response(&mut self) -> Option<Response> {
+        std::mem::replace(&mut self.response, None)
     }
 
     pub fn response(&self) -> &Option<Response> {
@@ -344,16 +380,26 @@ impl ResponseBuilder {
         self
     }
 
-    pub fn with_headers(mut self, headers: Vec<(HeaderName, &'static str)>) -> Self {
+    pub fn with_headers(mut self, headers: &[(HeaderName, &'static str)]) -> Self {
         self.response = self.response.set_headers(headers);
         self
     }
 
-    pub fn with_headers_str(mut self, headers: Vec<(&'static str, &'static str)>) -> Self {
+    pub fn with_headers_str(mut self, headers: &[(&'static str, &'static str)]) -> Self {
         self.response = self.response.set_headers_str(headers);
         self
     }
 
+    pub fn with_cookie(mut self, cookie: Cookie<'static>) -> Self {
+        self.response = self.response.set_cookie(cookie);
+        self
+    }
+
+    pub fn with_cookies(mut self, cookies: &[Cookie<'static>]) -> Self {
+        self.response = self.response.set_cookies(cookies);
+        self
+    }
+
     pub fn ok(mut self) -> ContextResult {
         *self.ctx.response_mut() = Some(self.response);
         Ok(self.ctx)

diff --git a/src/context/cookies.rs b/src/context/cookies.rs
@@ -0,0 +1,22 @@
+use cookie::CookieJar;
+
+#[derive(Default)]
+pub struct CookieParserData {
+    cookie_jar: CookieJar,
+}
+
+impl CookieParserData {
+    pub fn new() -> Self {
+        CookieParserData {
+            cookie_jar: CookieJar::new(),
+        }
+    }
+
+    pub fn cookie_jar(&self) -> &CookieJar {
+        &self.cookie_jar
+    }
+
+    pub fn cookie_jar_mut(&mut self) -> &mut CookieJar {
+        &mut self.cookie_jar
+    }
+}
diff --git a/src/context/memory_session.rs b/src/context/memory_session.rs
@@ -0,0 +1,33 @@
+use super::session::{SessionData, SessionStorage};
+
+use std::collections::HashMap;
+
+use lazy_static::lazy_static;
+use std::sync::Mutex;
+
+lazy_static! {
+    static ref STORAGE: Mutex<HashMap<String, SessionData>> = Mutex::new(HashMap::default());
+}
+
+#[derive(Default)]
+pub struct MemorySessionStorage {}
+
+impl MemorySessionStorage {
+    pub fn new() -> Self {
+        MemorySessionStorage {}
+    }
+}
+
+impl SessionStorage for MemorySessionStorage {
+    fn get_session(&self, id: &str) -> Option<SessionData> {
+        STORAGE.lock().unwrap().get(id).map(|x| (*x).clone())
+    }
+
+    fn set_session(&self, id: &str, session: SessionData) {
+        STORAGE.lock().unwrap().insert(id.to_string(), session);
+    }
+
+    fn remove_session(&self, id: &str) {
+        STORAGE.lock().unwrap().remove(id);
+    }
+}
diff --git a/src/context/session.rs b/src/context/session.rs
@@ -0,0 +1,62 @@
+use chrono::{DateTime, Duration, Local};
+
+use std::collections::HashMap;
+use std::ops::Add;
+
+pub trait SessionStorage: Send + Sync {
+    fn get_session(&self, id: &str) -> Option<SessionData>;
+    fn set_session(&self, id: &str, session: SessionData);
+    fn remove_session(&self, id: &str);
+}
+
+/// Session data to be consumed by context
+#[derive(Clone, Default)]
+pub struct SessionData {
+    data: HashMap<String, String>,
+    max_age: Option<DateTime<Local>>,
+    timeout: Option<Duration>,
+}
+
+impl SessionData {
+    pub fn new() -> Self {
+        SessionData {
+            data: HashMap::default(),
+            max_age: None,
+            timeout: None,
+        }
+    }
+
+    pub fn get(&self, name: &str) -> Option<&String> {
+        self.data.get(name)
+    }
+
+    pub fn get_all(&self) -> &HashMap<String, String> {
+        &self.data
+    }
+
+    pub fn set(&mut self, name: &str, value: impl ToString) -> Option<String> {
+        self.data.insert(name.to_string(), value.to_string())
+    }
+
+    pub fn refresh_session(&mut self) {
+        if let Some(timeout) = self.timeout {
+            self.max_age = Some(Local::now().add(timeout));
+        }
+    }
+
+    pub fn timeout(mut self, timeout: Duration) -> Self {
+        self.max_age = Some(Local::now().add(timeout));
+        self
+    }
+
+    pub fn is_alive(&self) -> bool {
+        match self.max_age {
+            Some(max_age) => max_age < Local::now(),
+            None => true,
+        }
+    }
+
+    pub fn remove(&mut self, name: &str) -> Option<String> {
+        self.data.remove(name)
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -10,4 +10,5 @@ pub mod router;
 pub use app::{App, EndpointExecutor};
 pub use error::ObsidianError;
 pub use hyper::{header, Body, HeaderMap, Method, Request, Response, StatusCode, Uri, Version};
+pub use cookie;
 pub use router::ContextResult;
diff --git a/src/middleware.rs b/src/middleware.rs
@@ -1,3 +1,5 @@
+pub mod cookie_parser;
+pub mod cookie_session;
 pub mod logger;
 
 use async_trait::async_trait;