feat: revamp cookies parser flow

Revamp the flow to fit current routing handler. Cookie flow and convertion are improved in this commit too.
obsidian-rs · May 3, 2020 · edf1fce · edf1fce
1 parent b70bbaa
commit edf1fce
Show file tree

Hide file tree

Showing 8 changed files with 102 additions and 31 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -18,6 +18,10 @@ path = 'examples/json.rs'
 name = 'app_state'
 path = 'examples/app_state.rs'
 
+[[example]]
+name = 'cookies'
+path = 'examples/cookies.rs'
+
 [package]
 name = 'obsidian'
 version = '0.2.2'

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
@@ -110,8 +110,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()
@@ -130,8 +130,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/src/app.rs b/src/app.rs
@@ -166,11 +166,33 @@ impl AppServer {
                             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
@@ -279,9 +279,7 @@ impl Context {
 
     pub fn cookie(&self, name: &str) -> Option<&Cookie> {
         if let Some(cookie_data) = self.get::<CookieParserData>() {
-            if let Some(ref cookie) = cookie_data.cookie_jar().get(name) {
-                return Some(*cookie);
-            }
+            return cookie_data.cookie_jar().get(name);
         }
 
         None
@@ -371,16 +369,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/middleware/cookie_parser.rs b/src/middleware/cookie_parser.rs
@@ -5,7 +5,7 @@ use hyper::header;
 use crate::app::EndpointExecutor;
 use crate::context::Context;
 use crate::middleware::Middleware;
-use crate::{Body, Response};
+use crate::router::ContextResult;
 
 #[derive(Default)]
 pub struct CookieParserData {
@@ -43,19 +43,17 @@ impl Middleware for CookieParser {
         &'a self,
         mut context: Context,
         ep_executor: EndpointExecutor<'a>,
-    ) -> Response<Body> {
+    ) -> ContextResult {
         if let Some(cookie_header) = context.headers().get(header::COOKIE) {
             if let Ok(cookie_str) = cookie_header.to_str() {
                 let mut cookie_data = CookieParserData::new();
 
                 cookie_str
                     .split("; ")
-                    .map(|x| x.trim().splitn(2, '=').collect::<Vec<&str>>())
-                    .for_each(|x| {
-                        if let (2, Some(k), Some(v)) = (x.len(), x.first(), x.last()) {
-                            cookie_data
-                                .cookie_jar_mut()
-                                .add_original(Cookie::new((*k).to_string(), (*v).to_string()));
+                    .map(|cookie_str_part| Cookie::parse(cookie_str_part.to_owned()))
+                    .for_each(|cookie_result| {
+                        if let Ok(cookie) = cookie_result {
+                            cookie_data.cookie_jar_mut().add_original(cookie);
                         }
                     });
 

diff --git a/src/router/responder.rs b/src/router/responder.rs
@@ -1,7 +1,7 @@
 use super::Response;
 use super::ResponseBody;
-use hyper::{header, StatusCode};
 use cookie::Cookie;
+use hyper::{header, StatusCode};
 
 pub trait Responder {
     fn respond_to(self) -> Response;
@@ -19,14 +19,14 @@ pub trait Responder {
         Response::new(self).set_header(key, value)
     }
 
-    fn with_headers(self, headers: Vec<(header::HeaderName, &'static str)>) -> Response
+    fn with_headers(self, headers: &[(header::HeaderName, &'static str)]) -> Response
     where
         Self: Responder + ResponseBody + Sized,
     {
         Response::new(self).set_headers(headers)
     }
 
-    fn with_headers_str(self, headers: Vec<(&'static str, &'static str)>) -> Response
+    fn with_headers_str(self, headers: &[(&'static str, &'static str)]) -> Response
     where
         Self: Responder + ResponseBody + Sized,
     {
@@ -40,7 +40,7 @@ pub trait Responder {
         Response::new(self).set_cookie(cookie)
     }
 
-    fn with_cookies(self, cookies: Vec<Cookie<'static>>) -> Response
+    fn with_cookies(self, cookies: &[Cookie<'static>]) -> Response
     where
         Self: Responder + ResponseBody + Sized,
     {

diff --git a/src/router/response.rs b/src/router/response.rs
@@ -44,6 +44,10 @@ impl Response {
         &self.cookies
     }
 
+    pub fn cookies_mut(&mut self) -> &mut Option<Vec<Cookie<'static>>> {
+        &mut self.cookies
+    }
+
     pub fn headers_mut(&mut self) -> &mut Option<Vec<(header::HeaderName, String)>> {
         &mut self.headers
     }
@@ -95,7 +99,7 @@ impl Response {
         self.set_cookie(cookie)
     }
 
-    pub fn with_cookies(self, cookies: Vec<Cookie<'static>>) -> Self {
+    pub fn with_cookies(self, cookies: &[Cookie<'static>]) -> Self {
         self.set_cookies(cookies)
     }
 
@@ -111,15 +115,15 @@ impl Response {
         self
     }
 
-    pub fn set_cookies(mut self, mut cookies: Vec<Cookie<'static>>) -> Self {
+    pub fn set_cookies(mut self, mut cookies: &[Cookie<'static>]) -> Self {
         match self.cookies {
-            Some(ref mut x) => x.append(&mut cookies),
-            None => self.cookies = Some(cookies),
+            Some(ref mut x) => x.extend_from_slice(&mut cookies),
+            None => self.cookies = Some(cookies.to_vec()),
         }
         self
     }
 
-    pub fn set_headers(mut self, headers: Vec<(header::HeaderName, &str)>) -> Self {
+    pub fn set_headers(mut self, headers: &[(header::HeaderName, &str)]) -> Self {
         let headers: Vec<(header::HeaderName, String)> = headers
             .iter()
             .map(|(k, v)| (header::HeaderName::from(k), (*v).to_string()))
@@ -133,11 +137,11 @@ impl Response {
     }
 
     // Alias set_headers method
-    pub fn with_headers(self, headers: Vec<(header::HeaderName, &'static str)>) -> Self {
+    pub fn with_headers(self, headers: &[(header::HeaderName, &'static str)]) -> Self {
         self.set_headers(headers)
     }
 
-    pub fn set_headers_str(mut self, headers: Vec<(&'static str, &'static str)>) -> Self {
+    pub fn set_headers_str(mut self, headers: &[(&'static str, &'static str)]) -> Self {
         let values: Vec<(header::HeaderName, String)> = headers
             .iter()
             .map(|(k, v)| {
@@ -156,7 +160,7 @@ impl Response {
     }
 
     // Alias set_headers_str method
-    pub fn with_headers_str(self, headers: Vec<(&'static str, &'static str)>) -> Self {
+    pub fn with_headers_str(self, headers: &[(&'static str, &'static str)]) -> Self {
         self.set_headers_str(headers)
     }